'''
model ensemble for cifar10  // input size(32,32)
'''


import torch
import torchvision

import copy
import torch.nn as nn
from resnetv2 import ResNet50 as resnet50v2


def split_resnet50(model):
    return nn.Sequential(
        model.conv1,
        model.layer1,
        model.layer2,
        model.layer3
    )


class PretrainEnsembleModel(nn.Module):

    def __init__(self):

        super(PretrainEnsembleModel, self).__init__()

        self.blocks = split_resnet50(resnet50v2())
        self.layer4_rotation = resnet50v2().layer4 
        self.layer4_jigsaw = resnet50v2().layer4

        self.fc_rotation = nn.Linear(2048, 4)
        self.fc_jigsaw = nn.Linear(2048, 31)

        self.avgpool1 = nn.AdaptiveAvgPool2d((1,1))
        self.avgpool2 = nn.AdaptiveAvgPool2d((1,1))
        self.avgpool3 = nn.AdaptiveAvgPool2d((1,1))


    def _Normal(self,x):
        mean=torch.Tensor([0.485, 0.456, 0.406])
        mean=mean[None,:,None,None].cuda()
        std = torch.Tensor([0.229, 0.224, 0.225])
        std = std[None,:,None,None].cuda()
        return x.sub(mean).div(std)

    def forward(self, x):

        feature_map = self.blocks(self._Normal(x))

        return feature_map

[((15, 11, 20, 5), 'torch.nn.Sequential', 'nn.Sequential', ({(16, 8, 16, 19): 'model.conv1', (17, 8, 17, 20): 'model.layer1', (18, 8, 18, 20): 'model.layer2', (19, 8, 19, 20): 'model.layer3'}, {}), '(model.conv1, model.layer1, model.layer2, model.layer3)', True, 'import torch.nn as nn\n'), ((33, 27, 33, 45), 'torch.nn.Linear', 'nn.Linear', ({(33, 37, 33, 41): '2048', (33, 43, 33, 44): '4'}, {}), '(2048, 4)', True, 'import torch.nn as nn\n'), ((34, 25, 34, 44), 'torch.nn.Linear', 'nn.Linear', ({(34, 35, 34, 39): '2048', (34, 41, 34, 43): '31'}, {}), '(2048, 31)', True, 'import torch.nn as nn\n'), ((36, 24, 36, 51), 'torch.nn.AdaptiveAvgPool2d', 'nn.AdaptiveAvgPool2d', ({(36, 45, 36, 50): '(1, 1)'}, {}), '((1, 1))', True, 'import torch.nn as nn\n'), ((37, 24, 37, 51), 'torch.nn.AdaptiveAvgPool2d', 'nn.AdaptiveAvgPool2d', ({(37, 45, 37, 50): '(1, 1)'}, {}), '((1, 1))', True, 'import torch.nn as nn\n'), ((38, 24, 38, 51), 'torch.nn.AdaptiveAvgPool2d', 'nn.AdaptiveAvgPool2d', ({(38, 45, 38, 50): '(1, 1)'}, {}), '((1, 1))', True, 'import torch.nn as nn\n'), ((42, 13, 42, 48), 'torch.Tensor', 'torch.Tensor', ({(42, 26, 42, 47): '[0.485, 0.456, 0.406]'}, {}), '([0.485, 0.456, 0.406])', False, 'import torch\n'), ((44, 14, 44, 49), 'torch.Tensor', 'torch.Tensor', ({(44, 27, 44, 48): '[0.229, 0.224, 0.225]'}, {}), '([0.229, 0.224, 0.225])', False, 'import torch\n'), ((29, 37, 29, 49), 'resnetv2.ResNet50', 'resnet50v2', ({}, {}), '()', True, 'from resnetv2 import ResNet50 as resnet50v2\n'), ((30, 31, 30, 43), 'resnetv2.ResNet50', 'resnet50v2', ({}, {}), '()', True, 'from resnetv2 import ResNet50 as resnet50v2\n'), ((31, 29, 31, 41), 'resnetv2.ResNet50', 'resnet50v2', ({}, {}), '()', True, 'from resnetv2 import ResNet50 as resnet50v2\n')]