import torch
import torch.nn as nn
from torch.nn import init
from torchvision import models
from torch.autograd import Variable
import torch.utils.model_zoo as model_zoo

import math
from collections import OrderedDict
import re
import torch.nn.functional as F

__all__ = ['resnet18', 'resnet50', 'densenet121']

model_urls = {
    'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
    'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
    'densenet121': 'https://download.pytorch.org/models/densenet121-a639ec97.pth',
    'wide_resnet50_2': 'https://download.pytorch.org/models/wide_resnet50_2-95faca4d.pth'
}

######################################################################
def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
    """3x3 convolution with padding"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False)

def conv1x1(in_planes, out_planes, stride=1):
    """1x1 convolution"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)

######################################################################
def weights_init_kaiming(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
        if m.bias is not None:
            init.constant_(m.bias.data, 0.0)        
    elif classname.find('Linear') != -1:
        init.kaiming_normal_(m.weight.data, a=0, mode='fan_out')
        if m.bias is not None:
            init.constant_(m.bias.data, 0.0)
    elif classname.find('BatchNorm1d') != -1:
        init.normal_(m.weight.data, 1.0, 0.02)
        if m.bias is not None:        
            init.constant_(m.bias.data, 0.0)
    elif classname.find('BatchNorm2d') != -1:
        init.normal_(m.weight.data, 1.0, 0.02)
        if m.bias is not None:        
            init.constant_(m.bias.data, 0.0)
    elif classname.find('BatchNorm') != -1:
        if m.affine is not None:
            init.constant_(m.weight.data, 1.0)
            init.constant_(m.bias.data, 0.0)        

def weights_init_normal(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        init.normal_(m.weight.data, 0.0, 0.02)
    elif classname.find('Linear') != -1:
        init.normal_(m.weight.data, 0.0, 0.02)
    elif classname.find('BatchNorm') != -1:
        init.normal_(m.weight.data, 1.0, 0.02)
        init.constant_(m.bias.data, 0.0)

def weights_init_xavier(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        init.xavier_normal_(m.weight.data, gain=1)
    elif classname.find('Linear') != -1:
        init.xavier_normal_(m.weight.data, gain=1)
    elif classname.find('BatchNorm') != -1:
        init.normal_(m.weight.data, 1.0, 0.02)
        init.constant_(m.bias.data, 0.0)

def weights_init_classifier(m):
    classname = m.__class__.__name__
    if classname.find('Linear') != -1:
        init.normal_(m.weight.data, std=0.001)
        init.constant_(m.bias.data, 0.0)     
        
######################################################################   
class ClassBlock(nn.Module):
    def __init__(self, input_dim, class_num, dropout=True, relu=True, num_bottleneck=512): #512
        super(ClassBlock, self).__init__()
        add_block = []
        add_block += [nn.Linear(input_dim, num_bottleneck)] 
        add_block += [nn.BatchNorm1d(num_bottleneck)]
        if relu: 
            add_block += [nn.ReLU()]
        if dropout: 
            add_block += [nn.Dropout(p=0.5)] 
        add_block = nn.Sequential(*add_block)
        add_block.apply(weights_init_kaiming)
        classifier = []
        classifier += [nn.Linear(num_bottleneck, class_num)]
        classifier = nn.Sequential(*classifier)
        classifier.apply(weights_init_classifier)
        self.add_block = add_block
        self.classifier = classifier
    def forward(self, x):
        x = self.add_block(x)
        x = self.classifier(x)
        return x    
    
######################################################################       
# Define the ResNet18-based Model
######################################################################     
class Res18_basic(nn.Module):
    def __init__(self, class_num):
        super(Res18_basic, self).__init__()
        fea_dim = 256
        model_ft = models.resnet18(pretrained=True)
        model_ft.avgpool = nn.AdaptiveAvgPool2d((1,1))
        model_ft.fc = nn.Sequential()
        self.model = model_ft
        self.fc_embed = nn.Linear(512, fea_dim)
        self.fc_embed.apply(weights_init_classifier)
        self.classifier = ClassBlock(512, class_num)
        self.classifier.apply(weights_init_classifier)
        
    def forward(self, x):
        x = self.model.conv1(x)
        x = self.model.bn1(x)
        x = self.model.relu(x)
        x = self.model.maxpool(x)
        x = self.model.layer1(x)
        x = self.model.layer2(x)
        x = self.model.layer3(x)
        x = self.model.layer4(x)
        x = self.model.avgpool(x)
        fea =  x.view(x.size(0), -1)
        embed_fea = self.fc_embed(fea)
        pred = self.classifier(fea)
        return embed_fea, pred 

# Define the ResNet18-based Model
class Res18(nn.Module):
    def __init__(self, class_num):
        super(Res18, self).__init__()
        fea_dim = 256
        model_ft = models.resnet18(pretrained=True)
        model_ft.avgpool = nn.AdaptiveAvgPool2d((1,1))
        model_ft.fc = nn.Sequential()
        self.model = model_ft
        self.fc_embed = nn.Linear(512, fea_dim)
        self.fc_embed.apply(weights_init_classifier)
        self.classifier = ClassBlock(512, class_num)
        self.classifier.apply(weights_init_classifier)
        
    def forward(self, x):
        x = self.model.conv1(x)
        x = self.model.bn1(x)
        x = self.model.relu(x)
        x = self.model.maxpool(x)
        x = self.model.layer1(x)
        x = self.model.layer2(x)
        x = self.model.layer3(x)
        x = self.model.layer4(x)
        x = self.model.avgpool(x)
        fea =  x.view(x.size(0), -1)
        embed_fea = self.fc_embed(fea)
        pred = self.classifier(fea)
        return embed_fea, pred 
    
        
class Res50(nn.Module):
    def __init__(self, class_num):
        super(Res50, self).__init__()
        fea_dim = 256        
        model_ft = models.resnet50(pretrained=True)
        model_ft.avgpool = nn.AdaptiveAvgPool2d((1,1))
        model_ft.fc = nn.Sequential()        
        self.model = model_ft
        self.fc_embed = nn.Linear(2048, fea_dim)
        self.fc_embed.apply(weights_init_classifier)
        self.classifier = ClassBlock(2048, class_num)
        self.classifier.apply(weights_init_classifier)        
        
    def forward(self, x):
        #print("x: ", x.size())
        x = self.model.conv1(x)
        #print("conv1(x): ", x.size())
        x = self.model.bn1(x)
        #print("bn1(x): ", x.size())
        x = self.model.relu(x)
        #print("relu(x): ", x.size())
        x = self.model.maxpool(x)
        #print("maxpool(x): ", x.size())
        x = self.model.layer1(x)
        #print("layer1(x): ", x.size())
        x = self.model.layer2(x)
        #print("layer2(x): ", x.size())
        x = self.model.layer3(x)
        #print("layer3(x): ", x.size())
        x = self.model.layer4(x)
        #print("layer4(x): ", x.size())
        x = self.model.avgpool(x)
        #print("avgpool(x): ", x.size())
        fea =  x.view(x.size(0), -1)
        embed_fea = self.fc_embed(fea)
        pred = self.classifier(fea)
        return embed_fea, pred     
        




class BasicBlock(nn.Module):
    def __init__(self, in_planes, out_planes, stride, dropRate=0.0, activate_before_residual=False):
        super(BasicBlock, self).__init__()
        self.bn1 = nn.BatchNorm2d(in_planes, momentum=0.001)
        self.relu1 = nn.LeakyReLU(negative_slope=0.1, inplace=True)
        self.conv1 = nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
                               padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_planes, momentum=0.001)
        self.relu2 = nn.LeakyReLU(negative_slope=0.1, inplace=True)
        self.conv2 = nn.Conv2d(out_planes, out_planes, kernel_size=3, stride=1,
                               padding=1, bias=False)
        self.droprate = dropRate
        self.equalInOut = (in_planes == out_planes)
        self.convShortcut = (not self.equalInOut) and nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride,
                               padding=0, bias=False) or None
        self.activate_before_residual = activate_before_residual
    def forward(self, x):
        if not self.equalInOut and self.activate_before_residual == True:
            x = self.relu1(self.bn1(x))
        else:
            out = self.relu1(self.bn1(x))
        out = self.relu2(self.bn2(self.conv1(out if self.equalInOut else x)))
        if self.droprate > 0:
            out = F.dropout(out, p=self.droprate, training=self.training)
        out = self.conv2(out)
        return torch.add(x if self.equalInOut else self.convShortcut(x), out)

class NetworkBlock(nn.Module):
    def __init__(self, nb_layers, in_planes, out_planes, block, stride, dropRate=0.0, activate_before_residual=False):
        super(NetworkBlock, self).__init__()
        self.layer = self._make_layer(block, in_planes, out_planes, nb_layers, stride, dropRate, activate_before_residual)
    def _make_layer(self, block, in_planes, out_planes, nb_layers, stride, dropRate, activate_before_residual):
        layers = []
        for i in range(int(nb_layers)):
            layers.append(block(i == 0 and in_planes or out_planes, out_planes, i == 0 and stride or 1, dropRate, activate_before_residual))
        return nn.Sequential(*layers)
    def forward(self, x):
        return self.layer(x)

class WideResNet(nn.Module):
    def __init__(self, num_classes, depth=28, widen_factor=2, dropRate=0.0):
        super(WideResNet, self).__init__()
        nChannels = [16, 16*widen_factor, 32*widen_factor, 64*widen_factor, 128*widen_factor]
        assert((depth - 4) % 6 == 0)
        n = (depth - 4) / 6
        block = BasicBlock
        # 1st conv before any network block
        self.conv1 = nn.Conv2d(3, nChannels[0], kernel_size=3, stride=2,
                               padding=1, bias=False)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        # 1st block
        self.block1 = NetworkBlock(n, nChannels[0], nChannels[1], block, 1, dropRate, activate_before_residual=True)
        # 2nd block
        self.block2 = NetworkBlock(n, nChannels[1], nChannels[2], block, 2, dropRate)
        # 3rd block
        self.block3 = NetworkBlock(n, nChannels[2], nChannels[3], block, 2, dropRate)
        # 4rd block
        self.block4 = NetworkBlock(n, nChannels[3], nChannels[4], block, 2, dropRate)
        # global average pooling and classifier
        self.bn1 = nn.BatchNorm2d(nChannels[4], momentum=0.001)
        self.relu = nn.LeakyReLU(negative_slope=0.1, inplace=True)
        self.fc = nn.Linear(nChannels[4], num_classes)
        self.nChannels = nChannels[4]
        self.classifier = ClassBlock(self.nChannels, num_classes)
        self.classifier.apply(weights_init_classifier) 

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
            elif isinstance(m, nn.Linear):
                nn.init.xavier_normal_(m.weight.data)
                m.bias.data.zero_()

    def forward(self, x):
        out = self.conv1(x)
        out = self.maxpool(out)
        out = self.block1(out)
        out = self.block2(out)
        out = self.block3(out)
        out = self.block4(out)
        out = self.relu(self.bn1(out))
        out = F.avg_pool2d(out, 7)
        out = out.view(-1, self.nChannels)
        #out = self.fc(out)
        out = self.classifier(out)
        return -1, out





'''
class Dense121(nn.Module):
    def __init__(self, class_num):
        super(Dense121, self).__init__()
        fea_dim = 256        
        model_ft = models.densenet121(pretrained=True)
        model_ft.features.classifier = nn.Sequential()
        model_ft.features.avgpool = nn.AdaptiveAvgPool2d((1,1))
        model_ft.features.fc_embed = nn.Linear(1024, fea_dim)
        model_ft.features.fc_embed.apply(weights_init_classifier)  
        model_ft.classifier = ClassBlock(2048, class_num)
        model_ft.classifier.apply(weights_init_classifier)  
        self.model = model_ft
        
    def forward(self, x):
        x = self.model.features.conv0(x)
        x = self.model.features.norm0(x)
        x = self.model.features.relu0(x)
        x = self.model.features.pool0(x)
        x = self.model.features.denseblock1(x)
        x = self.model.features.transition1(x)
        x = self.model.features.denseblock2(x)
        x = self.model.features.transition2(x)
        x = self.model.features.denseblock3(x)
        x = self.model.features.transition3(x)
        x = self.model.features.denseblock4(x)
        x = self.model.features.norm5(x)
        x = self.model.features.avgpool(x)
        fea =  x.view(x.size(0), -1)
        embed_fea = self.model.features.fc_embed(fea)
        pred = self.model.classifier(fea)        
        return embed_fea, pred              
        

class WideRes50_2(nn.Module):
    def __init__(self, class_num):
        super(WideRes50_2, self).__init__()
        fea_dim = 256        
        model_ft = wide_resnet50_2(pretrained=True)
        model_ft.avgpool = nn.AdaptiveAvgPool2d((1,1))
        model_ft.fc = nn.Sequential()        
        self.model = model_ft
        self.fc_embed = nn.Linear(2048, fea_dim)
        self.fc_embed.apply(weights_init_classifier)
        self.classifier = ClassBlock(2048, class_num)
        self.classifier.apply(weights_init_classifier)        
        
    def forward(self, x):
        x = self.model.conv1(x)
        x = self.model.bn1(x)
        x = self.model.relu(x)
        x = self.model.maxpool(x)
        x = self.model.layer1(x)
        x = self.model.layer2(x)
        x = self.model.layer3(x)
        x = self.model.layer4(x)
        x = self.model.avgpool(x)
        fea =  x.view(x.size(0), -1)
        embed_fea = self.fc_embed(fea)
        pred = self.classifier(fea)
        return embed_fea, pred

try:
    from torch.hub import load_state_dict_from_url
except ImportError:
    from torch.utils.model_zoo import load_url as load_state_dict_from_url

def wide_resnet50_2(pretrained=False, progress=True, **kwargs):
    
    kwargs['width_per_group'] = 64 * 2
    return _resnet('wide_resnet50_2', Bottleneck, [3, 4, 6, 3],
                   pretrained, progress, **kwargs)

class Bottleneck(nn.Module):
    # Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2)
    # while original implementation places the stride at the first 1x1 convolution(self.conv1)
    # according to "Deep residual learning for image recognition"https://arxiv.org/abs/1512.03385.
    # This variant is also known as ResNet V1.5 and improves accuracy according to
    # https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch.

    expansion = 4

    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
                 base_width=64, dilation=1, norm_layer=None):
        super(Bottleneck, self).__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        width = int(planes * (base_width / 64.)) * groups
        # Both self.conv2 and self.downsample layers downsample the input when stride != 1
        self.conv1 = conv1x1(inplanes, width)
        self.bn1 = norm_layer(width)
        self.conv2 = conv3x3(width, width, stride, groups, dilation)
        self.bn2 = norm_layer(width)
        self.conv3 = conv1x1(width, planes * self.expansion)
        self.bn3 = norm_layer(planes * self.expansion)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        identity = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.bn3(out)

        if self.downsample is not None:
            identity = self.downsample(x)

        out += identity
        out = self.relu(out)

        return out


class ResNet(nn.Module):

    def __init__(self, block, layers, num_classes=1000, zero_init_residual=False,
                 groups=1, width_per_group=64, replace_stride_with_dilation=None,
                 norm_layer=None):
        super(ResNet, self).__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        self._norm_layer = norm_layer

        self.inplanes = 64
        self.dilation = 1
        if replace_stride_with_dilation is None:
            # each element in the tuple indicates if we should replace
            # the 2x2 stride with a dilated convolution instead
            replace_stride_with_dilation = [False, False, False]
        if len(replace_stride_with_dilation) != 3:
            raise ValueError("replace_stride_with_dilation should be None "
                             "or a 3-element tuple, got {}".format(replace_stride_with_dilation))
        self.groups = groups
        self.base_width = width_per_group
        self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3,
                               bias=False)
        self.bn1 = norm_layer(self.inplanes)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        self.layer1 = self._make_layer(block, 64, layers[0])
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2,
                                       dilate=replace_stride_with_dilation[0])
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2,
                                       dilate=replace_stride_with_dilation[1])
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2,
                                       dilate=replace_stride_with_dilation[2])
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(512 * block.expansion, num_classes)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

        # Zero-initialize the last BN in each residual branch,
        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
        if zero_init_residual:
            for m in self.modules():
                if isinstance(m, Bottleneck):
                    nn.init.constant_(m.bn3.weight, 0)
                elif isinstance(m, BasicBlock):
                    nn.init.constant_(m.bn2.weight, 0)

    def _make_layer(self, block, planes, blocks, stride=1, dilate=False):
        norm_layer = self._norm_layer
        downsample = None
        previous_dilation = self.dilation
        if dilate:
            self.dilation *= stride
            stride = 1
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                conv1x1(self.inplanes, planes * block.expansion, stride),
                norm_layer(planes * block.expansion),
            )

        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample, self.groups,
                            self.base_width, previous_dilation, norm_layer))
        self.inplanes = planes * block.expansion
        for _ in range(1, blocks):
            layers.append(block(self.inplanes, planes, groups=self.groups,
                                base_width=self.base_width, dilation=self.dilation,
                                norm_layer=norm_layer))

        return nn.Sequential(*layers)

    def _forward_impl(self, x):
        # See note [TorchScript super()]
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.fc(x)

        return x

    def forward(self, x):
        return self._forward_impl(x)


def _resnet(arch, block, layers, pretrained, progress, **kwargs):
    model = ResNet(block, layers, **kwargs)
    if pretrained:
        state_dict = load_state_dict_from_url(model_urls[arch],
                                              progress=progress)
        model.load_state_dict(state_dict)
    return model
'''

import torch
import torch.nn as nn
from torch.nn import init
from torchvision import models
from torch.autograd import Variable
import torch.utils.model_zoo as model_zoo

import math
from collections import OrderedDict
import re
import torch.nn.functional as F

__all__ = ['resnet18', 'resnet50', 'densenet121']

model_urls = {
    'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
    'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
    'densenet121': 'https://download.pytorch.org/models/densenet121-a639ec97.pth',
    'wide_resnet50_2': 'https://download.pytorch.org/models/wide_resnet50_2-95faca4d.pth'
}

######################################################################
def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
    """3x3 convolution with padding"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False)

def conv1x1(in_planes, out_planes, stride=1):
    """1x1 convolution"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)

######################################################################
def weights_init_kaiming(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
        if m.bias is not None:
            init.constant_(m.bias.data, 0.0)        
    elif classname.find('Linear') != -1:
        init.kaiming_normal_(m.weight.data, a=0, mode='fan_out')
        if m.bias is not None:
            init.constant_(m.bias.data, 0.0)
    elif classname.find('BatchNorm1d') != -1:
        init.normal_(m.weight.data, 1.0, 0.02)
        if m.bias is not None:        
            init.constant_(m.bias.data, 0.0)
    elif classname.find('BatchNorm2d') != -1:
        init.normal_(m.weight.data, 1.0, 0.02)
        if m.bias is not None:        
            init.constant_(m.bias.data, 0.0)
    elif classname.find('BatchNorm') != -1:
        if m.affine is not None:
            init.constant_(m.weight.data, 1.0)
            init.constant_(m.bias.data, 0.0)        

def weights_init_normal(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        init.normal_(m.weight.data, 0.0, 0.02)
    elif classname.find('Linear') != -1:
        init.normal_(m.weight.data, 0.0, 0.02)
    elif classname.find('BatchNorm') != -1:
        init.normal_(m.weight.data, 1.0, 0.02)
        init.constant_(m.bias.data, 0.0)

def weights_init_xavier(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        init.xavier_normal_(m.weight.data, gain=1)
    elif classname.find('Linear') != -1:
        init.xavier_normal_(m.weight.data, gain=1)
    elif classname.find('BatchNorm') != -1:
        init.normal_(m.weight.data, 1.0, 0.02)
        init.constant_(m.bias.data, 0.0)

def weights_init_classifier(m):
    classname = m.__class__.__name__
    if classname.find('Linear') != -1:
        init.normal_(m.weight.data, std=0.001)
        init.constant_(m.bias.data, 0.0)     
        
######################################################################   
class ClassBlock(nn.Module):
    def __init__(self, input_dim, class_num, dropout=True, relu=True, num_bottleneck=512): #512
        super(ClassBlock, self).__init__()
        add_block = []
        add_block += [nn.Linear(input_dim, num_bottleneck)] 
        add_block += [nn.BatchNorm1d(num_bottleneck)]
        if relu: 
            add_block += [nn.ReLU()]
        if dropout: 
            add_block += [nn.Dropout(p=0.5)] 
        add_block = nn.Sequential(*add_block)
        add_block.apply(weights_init_kaiming)
        classifier = []
        classifier += [nn.Linear(num_bottleneck, class_num)]
        classifier = nn.Sequential(*classifier)
        classifier.apply(weights_init_classifier)
        self.add_block = add_block
        self.classifier = classifier
    def forward(self, x):
        x = self.add_block(x)
        x = self.classifier(x)
        return x    
    
######################################################################       
# Define the ResNet18-based Model
######################################################################     
class Res18_basic(nn.Module):
    def __init__(self, class_num):
        super(Res18_basic, self).__init__()
        fea_dim = 256
        model_ft = models.resnet18(pretrained=True)
        model_ft.avgpool = nn.AdaptiveAvgPool2d((1,1))
        model_ft.fc = nn.Sequential()
        self.model = model_ft
        self.fc_embed = nn.Linear(512, fea_dim)
        self.fc_embed.apply(weights_init_classifier)
        self.classifier = ClassBlock(512, class_num)
        self.classifier.apply(weights_init_classifier)
        
    def forward(self, x):
        x = self.model.conv1(x)
        x = self.model.bn1(x)
        x = self.model.relu(x)
        x = self.model.maxpool(x)
        x = self.model.layer1(x)
        x = self.model.layer2(x)
        x = self.model.layer3(x)
        x = self.model.layer4(x)
        x = self.model.avgpool(x)
        fea =  x.view(x.size(0), -1)
        embed_fea = self.fc_embed(fea)
        pred = self.classifier(fea)
        return embed_fea, pred 

# Define the ResNet18-based Model
class Res18(nn.Module):
    def __init__(self, class_num):
        super(Res18, self).__init__()
        fea_dim = 256
        model_ft = models.resnet18(pretrained=True)
        model_ft.avgpool = nn.AdaptiveAvgPool2d((1,1))
        model_ft.fc = nn.Sequential()
        self.model = model_ft
        self.fc_embed = nn.Linear(512, fea_dim)
        self.fc_embed.apply(weights_init_classifier)
        self.classifier = ClassBlock(512, class_num)
        self.classifier.apply(weights_init_classifier)
        
    def forward(self, x):
        x = self.model.conv1(x)
        x = self.model.bn1(x)
        x = self.model.relu(x)
        x = self.model.maxpool(x)
        x = self.model.layer1(x)
        x = self.model.layer2(x)
        x = self.model.layer3(x)
        x = self.model.layer4(x)
        x = self.model.avgpool(x)
        fea =  x.view(x.size(0), -1)
        embed_fea = self.fc_embed(fea)
        pred = self.classifier(fea)
        return embed_fea, pred 
    
        
class Res50(nn.Module):
    def __init__(self, class_num):
        super(Res50, self).__init__()
        fea_dim = 256        
        model_ft = models.resnet50(pretrained=True)
        model_ft.avgpool = nn.AdaptiveAvgPool2d((1,1))
        model_ft.fc = nn.Sequential()        
        self.model = model_ft
        self.fc_embed = nn.Linear(2048, fea_dim)
        self.fc_embed.apply(weights_init_classifier)
        self.classifier = ClassBlock(2048, class_num)
        self.classifier.apply(weights_init_classifier)        
        
    def forward(self, x):
        #print("x: ", x.size())
        x = self.model.conv1(x)
        #print("conv1(x): ", x.size())
        x = self.model.bn1(x)
        #print("bn1(x): ", x.size())
        x = self.model.relu(x)
        #print("relu(x): ", x.size())
        x = self.model.maxpool(x)
        #print("maxpool(x): ", x.size())
        x = self.model.layer1(x)
        #print("layer1(x): ", x.size())
        x = self.model.layer2(x)
        #print("layer2(x): ", x.size())
        x = self.model.layer3(x)
        #print("layer3(x): ", x.size())
        x = self.model.layer4(x)
        #print("layer4(x): ", x.size())
        x = self.model.avgpool(x)
        #print("avgpool(x): ", x.size())
        fea =  x.view(x.size(0), -1)
        embed_fea = self.fc_embed(fea)
        pred = self.classifier(fea)
        return embed_fea, pred     
        




class BasicBlock(nn.Module):
    def __init__(self, in_planes, out_planes, stride, dropRate=0.0, activate_before_residual=False):
        super(BasicBlock, self).__init__()
        self.bn1 = nn.BatchNorm2d(in_planes, momentum=0.001)
        self.relu1 = nn.LeakyReLU(negative_slope=0.1, inplace=True)
        self.conv1 = nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
                               padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_planes, momentum=0.001)
        self.relu2 = nn.LeakyReLU(negative_slope=0.1, inplace=True)
        self.conv2 = nn.Conv2d(out_planes, out_planes, kernel_size=3, stride=1,
                               padding=1, bias=False)
        self.droprate = dropRate
        self.equalInOut = (in_planes == out_planes)
        self.convShortcut = (not self.equalInOut) and nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride,
                               padding=0, bias=False) or None
        self.activate_before_residual = activate_before_residual
    def forward(self, x):
        if not self.equalInOut and self.activate_before_residual == True:
            x = self.relu1(self.bn1(x))
        else:
            out = self.relu1(self.bn1(x))
        out = self.relu2(self.bn2(self.conv1(out if self.equalInOut else x)))
        if self.droprate > 0:
            out = F.dropout(out, p=self.droprate, training=self.training)
        out = self.conv2(out)
        return torch.add(x if self.equalInOut else self.convShortcut(x), out)

class NetworkBlock(nn.Module):
    def __init__(self, nb_layers, in_planes, out_planes, block, stride, dropRate=0.0, activate_before_residual=False):
        super(NetworkBlock, self).__init__()
        self.layer = self._make_layer(block, in_planes, out_planes, nb_layers, stride, dropRate, activate_before_residual)
    def _make_layer(self, block, in_planes, out_planes, nb_layers, stride, dropRate, activate_before_residual):
        layers = []
        for i in range(int(nb_layers)):
            layers.append(block(i == 0 and in_planes or out_planes, out_planes, i == 0 and stride or 1, dropRate, activate_before_residual))
        return nn.Sequential(*layers)
    def forward(self, x):
        return self.layer(x)

class WideResNet(nn.Module):
    def __init__(self, num_classes, depth=28, widen_factor=2, dropRate=0.0):
        super(WideResNet, self).__init__()
        nChannels = [16, 16*widen_factor, 32*widen_factor, 64*widen_factor, 128*widen_factor]
        assert((depth - 4) % 6 == 0)
        n = (depth - 4) / 6
        block = BasicBlock
        # 1st conv before any network block
        self.conv1 = nn.Conv2d(3, nChannels[0], kernel_size=3, stride=2,
                               padding=1, bias=False)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        # 1st block
        self.block1 = NetworkBlock(n, nChannels[0], nChannels[1], block, 1, dropRate, activate_before_residual=True)
        # 2nd block
        self.block2 = NetworkBlock(n, nChannels[1], nChannels[2], block, 2, dropRate)
        # 3rd block
        self.block3 = NetworkBlock(n, nChannels[2], nChannels[3], block, 2, dropRate)
        # 4rd block
        self.block4 = NetworkBlock(n, nChannels[3], nChannels[4], block, 2, dropRate)
        # global average pooling and classifier
        self.bn1 = nn.BatchNorm2d(nChannels[4], momentum=0.001)
        self.relu = nn.LeakyReLU(negative_slope=0.1, inplace=True)
        self.fc = nn.Linear(nChannels[4], num_classes)
        self.nChannels = nChannels[4]
        self.classifier = ClassBlock(self.nChannels, num_classes)
        self.classifier.apply(weights_init_classifier) 

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
            elif isinstance(m, nn.Linear):
                nn.init.xavier_normal_(m.weight.data)
                m.bias.data.zero_()

    def forward(self, x):
        out = self.conv1(x)
        out = self.maxpool(out)
        out = self.block1(out)
        out = self.block2(out)
        out = self.block3(out)
        out = self.block4(out)
        out = self.relu(self.bn1(out))
        out = F.avg_pool2d(out, 7)
        out = out.view(-1, self.nChannels)
        #out = self.fc(out)
        out = self.classifier(out)
        return -1, out





'''
class Dense121(nn.Module):
    def __init__(self, class_num):
        super(Dense121, self).__init__()
        fea_dim = 256        
        model_ft = models.densenet121(pretrained=True)
        model_ft.features.classifier = nn.Sequential()
        model_ft.features.avgpool = nn.AdaptiveAvgPool2d((1,1))
        model_ft.features.fc_embed = nn.Linear(1024, fea_dim)
        model_ft.features.fc_embed.apply(weights_init_classifier)  
        model_ft.classifier = ClassBlock(2048, class_num)
        model_ft.classifier.apply(weights_init_classifier)  
        self.model = model_ft
        
    def forward(self, x):
        x = self.model.features.conv0(x)
        x = self.model.features.norm0(x)
        x = self.model.features.relu0(x)
        x = self.model.features.pool0(x)
        x = self.model.features.denseblock1(x)
        x = self.model.features.transition1(x)
        x = self.model.features.denseblock2(x)
        x = self.model.features.transition2(x)
        x = self.model.features.denseblock3(x)
        x = self.model.features.transition3(x)
        x = self.model.features.denseblock4(x)
        x = self.model.features.norm5(x)
        x = self.model.features.avgpool(x)
        fea =  x.view(x.size(0), -1)
        embed_fea = self.model.features.fc_embed(fea)
        pred = self.model.classifier(fea)        
        return embed_fea, pred              
        

class WideRes50_2(nn.Module):
    def __init__(self, class_num):
        super(WideRes50_2, self).__init__()
        fea_dim = 256        
        model_ft = wide_resnet50_2(pretrained=True)
        model_ft.avgpool = nn.AdaptiveAvgPool2d((1,1))
        model_ft.fc = nn.Sequential()        
        self.model = model_ft
        self.fc_embed = nn.Linear(2048, fea_dim)
        self.fc_embed.apply(weights_init_classifier)
        self.classifier = ClassBlock(2048, class_num)
        self.classifier.apply(weights_init_classifier)        
        
    def forward(self, x):
        x = self.model.conv1(x)
        x = self.model.bn1(x)
        x = self.model.relu(x)
        x = self.model.maxpool(x)
        x = self.model.layer1(x)
        x = self.model.layer2(x)
        x = self.model.layer3(x)
        x = self.model.layer4(x)
        x = self.model.avgpool(x)
        fea =  x.view(x.size(0), -1)
        embed_fea = self.fc_embed(fea)
        pred = self.classifier(fea)
        return embed_fea, pred

try:
    from torch.hub import load_state_dict_from_url
except ImportError:
    from torch.utils.model_zoo import load_url as load_state_dict_from_url

def wide_resnet50_2(pretrained=False, progress=True, **kwargs):
    
    kwargs['width_per_group'] = 64 * 2
    return _resnet('wide_resnet50_2', Bottleneck, [3, 4, 6, 3],
                   pretrained, progress, **kwargs)

class Bottleneck(nn.Module):
    # Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2)
    # while original implementation places the stride at the first 1x1 convolution(self.conv1)
    # according to "Deep residual learning for image recognition"https://arxiv.org/abs/1512.03385.
    # This variant is also known as ResNet V1.5 and improves accuracy according to
    # https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch.

    expansion = 4

    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
                 base_width=64, dilation=1, norm_layer=None):
        super(Bottleneck, self).__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        width = int(planes * (base_width / 64.)) * groups
        # Both self.conv2 and self.downsample layers downsample the input when stride != 1
        self.conv1 = conv1x1(inplanes, width)
        self.bn1 = norm_layer(width)
        self.conv2 = conv3x3(width, width, stride, groups, dilation)
        self.bn2 = norm_layer(width)
        self.conv3 = conv1x1(width, planes * self.expansion)
        self.bn3 = norm_layer(planes * self.expansion)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        identity = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.bn3(out)

        if self.downsample is not None:
            identity = self.downsample(x)

        out += identity
        out = self.relu(out)

        return out


class ResNet(nn.Module):

    def __init__(self, block, layers, num_classes=1000, zero_init_residual=False,
                 groups=1, width_per_group=64, replace_stride_with_dilation=None,
                 norm_layer=None):
        super(ResNet, self).__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        self._norm_layer = norm_layer

        self.inplanes = 64
        self.dilation = 1
        if replace_stride_with_dilation is None:
            # each element in the tuple indicates if we should replace
            # the 2x2 stride with a dilated convolution instead
            replace_stride_with_dilation = [False, False, False]
        if len(replace_stride_with_dilation) != 3:
            raise ValueError("replace_stride_with_dilation should be None "
                             "or a 3-element tuple, got {}".format(replace_stride_with_dilation))
        self.groups = groups
        self.base_width = width_per_group
        self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3,
                               bias=False)
        self.bn1 = norm_layer(self.inplanes)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        self.layer1 = self._make_layer(block, 64, layers[0])
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2,
                                       dilate=replace_stride_with_dilation[0])
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2,
                                       dilate=replace_stride_with_dilation[1])
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2,
                                       dilate=replace_stride_with_dilation[2])
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(512 * block.expansion, num_classes)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

        # Zero-initialize the last BN in each residual branch,
        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
        if zero_init_residual:
            for m in self.modules():
                if isinstance(m, Bottleneck):
                    nn.init.constant_(m.bn3.weight, 0)
                elif isinstance(m, BasicBlock):
                    nn.init.constant_(m.bn2.weight, 0)

    def _make_layer(self, block, planes, blocks, stride=1, dilate=False):
        norm_layer = self._norm_layer
        downsample = None
        previous_dilation = self.dilation
        if dilate:
            self.dilation *= stride
            stride = 1
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                conv1x1(self.inplanes, planes * block.expansion, stride),
                norm_layer(planes * block.expansion),
            )

        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample, self.groups,
                            self.base_width, previous_dilation, norm_layer))
        self.inplanes = planes * block.expansion
        for _ in range(1, blocks):
            layers.append(block(self.inplanes, planes, groups=self.groups,
                                base_width=self.base_width, dilation=self.dilation,
                                norm_layer=norm_layer))

        return nn.Sequential(*layers)

    def _forward_impl(self, x):
        # See note [TorchScript super()]
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.fc(x)

        return x

    def forward(self, x):
        return self._forward_impl(x)


def _resnet(arch, block, layers, pretrained, progress, **kwargs):
    model = ResNet(block, layers, **kwargs)
    if pretrained:
        state_dict = load_state_dict_from_url(model_urls[arch],
                                              progress=progress)
        model.load_state_dict(state_dict)
    return model
'''

###################################################################### 3x3 convolution with padding 1x1 convolution ###################################################################### ###################################################################### #512 ###################################################################### # Define the ResNet18-based Model ###################################################################### # Define the ResNet18-based Model #print("x: ", x.size()) #print("conv1(x): ", x.size()) #print("bn1(x): ", x.size()) #print("relu(x): ", x.size()) #print("maxpool(x): ", x.size()) #print("layer1(x): ", x.size()) #print("layer2(x): ", x.size()) #print("layer3(x): ", x.size()) #print("layer4(x): ", x.size()) #print("avgpool(x): ", x.size()) # 1st conv before any network block # 1st block # 2nd block # 3rd block # 4rd block # global average pooling and classifier #out = self.fc(out) class Dense121(nn.Module):     def __init__(self, class_num):         super(Dense121, self).__init__()         fea_dim = 256                 model_ft = models.densenet121(pretrained=True)         model_ft.features.classifier = nn.Sequential()         model_ft.features.avgpool = nn.AdaptiveAvgPool2d((1,1))         model_ft.features.fc_embed = nn.Linear(1024, fea_dim)         model_ft.features.fc_embed.apply(weights_init_classifier)           model_ft.classifier = ClassBlock(2048, class_num)         model_ft.classifier.apply(weights_init_classifier)           self.model = model_ft              def forward(self, x):         x = self.model.features.conv0(x)         x = self.model.features.norm0(x)         x = self.model.features.relu0(x)         x = self.model.features.pool0(x)         x = self.model.features.denseblock1(x)         x = self.model.features.transition1(x)         x = self.model.features.denseblock2(x)         x = self.model.features.transition2(x)         x = self.model.features.denseblock3(x)         x = self.model.features.transition3(x)         x = self.model.features.denseblock4(x)         x = self.model.features.norm5(x)         x = self.model.features.avgpool(x)         fea =  x.view(x.size(0), -1)         embed_fea = self.model.features.fc_embed(fea)         pred = self.model.classifier(fea)                 return embed_fea, pred                         class WideRes50_2(nn.Module):     def __init__(self, class_num):         super(WideRes50_2, self).__init__()         fea_dim = 256                 model_ft = wide_resnet50_2(pretrained=True)         model_ft.avgpool = nn.AdaptiveAvgPool2d((1,1))         model_ft.fc = nn.Sequential()                 self.model = model_ft         self.fc_embed = nn.Linear(2048, fea_dim)         self.fc_embed.apply(weights_init_classifier)         self.classifier = ClassBlock(2048, class_num)         self.classifier.apply(weights_init_classifier)                      def forward(self, x):         x = self.model.conv1(x)         x = self.model.bn1(x)         x = self.model.relu(x)         x = self.model.maxpool(x)         x = self.model.layer1(x)         x = self.model.layer2(x)         x = self.model.layer3(x)         x = self.model.layer4(x)         x = self.model.avgpool(x)         fea =  x.view(x.size(0), -1)         embed_fea = self.fc_embed(fea)         pred = self.classifier(fea)         return embed_fea, pred  try:     from torch.hub import load_state_dict_from_url except ImportError:     from torch.utils.model_zoo import load_url as load_state_dict_from_url  def wide_resnet50_2(pretrained=False, progress=True, **kwargs):          kwargs['width_per_group'] = 64 * 2     return _resnet('wide_resnet50_2', Bottleneck, [3, 4, 6, 3],                    pretrained, progress, **kwargs)  class Bottleneck(nn.Module):     # Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2)     # while original implementation places the stride at the first 1x1 convolution(self.conv1)     # according to "Deep residual learning for image recognition"https://arxiv.org/abs/1512.03385.     # This variant is also known as ResNet V1.5 and improves accuracy according to     # https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch.      expansion = 4      def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,                  base_width=64, dilation=1, norm_layer=None):         super(Bottleneck, self).__init__()         if norm_layer is None:             norm_layer = nn.BatchNorm2d         width = int(planes * (base_width / 64.)) * groups         # Both self.conv2 and self.downsample layers downsample the input when stride != 1         self.conv1 = conv1x1(inplanes, width)         self.bn1 = norm_layer(width)         self.conv2 = conv3x3(width, width, stride, groups, dilation)         self.bn2 = norm_layer(width)         self.conv3 = conv1x1(width, planes * self.expansion)         self.bn3 = norm_layer(planes * self.expansion)         self.relu = nn.ReLU(inplace=True)         self.downsample = downsample         self.stride = stride      def forward(self, x):         identity = x          out = self.conv1(x)         out = self.bn1(out)         out = self.relu(out)          out = self.conv2(out)         out = self.bn2(out)         out = self.relu(out)          out = self.conv3(out)         out = self.bn3(out)          if self.downsample is not None:             identity = self.downsample(x)          out += identity         out = self.relu(out)          return out   class ResNet(nn.Module):      def __init__(self, block, layers, num_classes=1000, zero_init_residual=False,                  groups=1, width_per_group=64, replace_stride_with_dilation=None,                  norm_layer=None):         super(ResNet, self).__init__()         if norm_layer is None:             norm_layer = nn.BatchNorm2d         self._norm_layer = norm_layer          self.inplanes = 64         self.dilation = 1         if replace_stride_with_dilation is None:             # each element in the tuple indicates if we should replace             # the 2x2 stride with a dilated convolution instead             replace_stride_with_dilation = [False, False, False]         if len(replace_stride_with_dilation) != 3:             raise ValueError("replace_stride_with_dilation should be None "                              "or a 3-element tuple, got {}".format(replace_stride_with_dilation))         self.groups = groups         self.base_width = width_per_group         self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3,                                bias=False)         self.bn1 = norm_layer(self.inplanes)         self.relu = nn.ReLU(inplace=True)         self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)         self.layer1 = self._make_layer(block, 64, layers[0])         self.layer2 = self._make_layer(block, 128, layers[1], stride=2,                                        dilate=replace_stride_with_dilation[0])         self.layer3 = self._make_layer(block, 256, layers[2], stride=2,                                        dilate=replace_stride_with_dilation[1])         self.layer4 = self._make_layer(block, 512, layers[3], stride=2,                                        dilate=replace_stride_with_dilation[2])         self.avgpool = nn.AdaptiveAvgPool2d((1, 1))         self.fc = nn.Linear(512 * block.expansion, num_classes)          for m in self.modules():             if isinstance(m, nn.Conv2d):                 nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')             elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):                 nn.init.constant_(m.weight, 1)                 nn.init.constant_(m.bias, 0)          # Zero-initialize the last BN in each residual branch,         # so that the residual branch starts with zeros, and each residual block behaves like an identity.         # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677         if zero_init_residual:             for m in self.modules():                 if isinstance(m, Bottleneck):                     nn.init.constant_(m.bn3.weight, 0)                 elif isinstance(m, BasicBlock):                     nn.init.constant_(m.bn2.weight, 0)      def _make_layer(self, block, planes, blocks, stride=1, dilate=False):         norm_layer = self._norm_layer         downsample = None         previous_dilation = self.dilation         if dilate:             self.dilation *= stride             stride = 1         if stride != 1 or self.inplanes != planes * block.expansion:             downsample = nn.Sequential(                 conv1x1(self.inplanes, planes * block.expansion, stride),                 norm_layer(planes * block.expansion),             )          layers = []         layers.append(block(self.inplanes, planes, stride, downsample, self.groups,                             self.base_width, previous_dilation, norm_layer))         self.inplanes = planes * block.expansion         for _ in range(1, blocks):             layers.append(block(self.inplanes, planes, groups=self.groups,                                 base_width=self.base_width, dilation=self.dilation,                                 norm_layer=norm_layer))          return nn.Sequential(*layers)      def _forward_impl(self, x):         # See note [TorchScript super()]         x = self.conv1(x)         x = self.bn1(x)         x = self.relu(x)         x = self.maxpool(x)          x = self.layer1(x)         x = self.layer2(x)         x = self.layer3(x)         x = self.layer4(x)          x = self.avgpool(x)         x = torch.flatten(x, 1)         x = self.fc(x)          return x      def forward(self, x):         return self._forward_impl(x)   def _resnet(arch, block, layers, pretrained, progress, **kwargs):     model = ResNet(block, layers, **kwargs)     if pretrained:         state_dict = load_state_dict_from_url(model_urls[arch],                                               progress=progress)         model.load_state_dict(state_dict)     return model