code: from __future__ import absolute_import
  
  import torch
  from torch import nn
  from torch.nn import functional as F
  from torch.nn import init
  import torchvision
  from collections import OrderedDict
  
  from ..models.layers.adain import SMMBlock
  
  
  __all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
             'resnet152', 'resnet50_mldg_smm']
  
  
  class ResNet(nn.Module):
      __factory = {
          18: torchvision.models.resnet18,
          34: torchvision.models.resnet34,
          50: torchvision.models.resnet50,
          101: torchvision.models.resnet101,
          152: torchvision.models.resnet152,
      }
  
      def __init__(self, depth, pretrained=True, cut_at_pooling=False,
                   num_features=0, norm=False, dropout=0, num_classes=None):
          super(ResNet, self).__init__()
          self.pretrained = pretrained
          self.depth = depth
          self.cut_at_pooling = cut_at_pooling
  
          # Construct base (pretrained) resnet
          if depth not in ResNet.__factory:
              raise KeyError("Unsupported depth:", depth)
          resnet = ResNet.__factory[depth](pretrained=pretrained)
          resnet.layer4[0].conv2.stride = (1,1)
          resnet.layer4[0].downsample[0].stride = (1,1)
          # self.base = nn.Sequential(
          #     resnet.conv1, resnet.bn1, resnet.relu, resnet.maxpool,
          #     resnet.layer1, resnet.layer2, resnet.layer3, resnet.layer4)
  
          self.conv = nn.Sequential(OrderedDict([
              ('conv1', resnet.conv1),
              ('bn1', resnet.bn1),
              ('relu', resnet.relu),
              ('maxpool', resnet.maxpool)]))
  
          self.layer1 = resnet.layer1
          self.layer2 = resnet.layer2
          self.layer3 = resnet.layer3
          self.layer4 = resnet.layer4
  
          self.gap = nn.AdaptiveAvgPool2d(1)
  
          self.smm_block = SMMBlock(1, rand=False, learnable=False)
  
          if not self.cut_at_pooling:
              self.num_features = num_features
              self.norm = norm
              self.dropout = dropout
              self.has_embedding = num_features > 0
              self.num_classes = num_classes
  
              out_planes = resnet.fc.in_features
  
              # Append new layers
              if self.has_embedding:
                  self.feat = nn.Linear(out_planes, self.num_features)
                  self.feat_bn = nn.BatchNorm1d(self.num_features)
                  init.kaiming_normal_(self.feat.weight, mode='fan_out')
                  init.constant_(self.feat.bias, 0)
              else:
                  # Change the num_features to CNN output channels
                  self.num_features = out_planes
                  self.feat_bn = nn.BatchNorm1d(self.num_features)
              self.feat_bn.bias.requires_grad_(False)
              if self.dropout > 0:
                  self.drop = nn.Dropout(self.dropout)
  
              self.classifier = nn.Linear(self.num_features, self.num_classes, bias=False)
              init.normal_(self.classifier.weight, std=0.001)
  
          init.constant_(self.feat_bn.weight, 1)
          init.constant_(self.feat_bn.bias, 0)
  
          if not pretrained:
              self.reset_params()
  
      def forward(self, x, meta_train=True, output_prob=False, return_featuremaps=False):
          if self.training:
              num_domains = len(x)
              x = torch.cat(x, dim=0)
  
          x = self.conv(x)
  
          # NOTE: change to 'if self.training and meta_train:'
          if meta_train:
              mixed_x, _ = self.smm_block(x)
              if return_featuremaps:
                  return [x, mixed_x]
              x = mixed_x
  
  
          x = self.layer1(x)
          x = self.layer2(x)
          x = self.layer3(x)
          x = self.layer4(x)
  
          x = self.gap(x)
          x = x.view(x.size(0), -1)
  
          if self.cut_at_pooling:
              return x
  
          if self.has_embedding:
              bn_x = self.feat_bn(self.feat(x))
          else:
              bn_x = self.feat_bn(x)
  
          if self.training is False and output_prob is False:
              bn_x = F.normalize(bn_x)
              return bn_x
  
          if self.norm:
              norm_bn_x = F.normalize(bn_x)
          elif self.has_embedding:
              bn_x = F.relu(bn_x)
  
          if self.dropout > 0:
              bn_x = self.drop(bn_x)
  
          prob = self.classifier(bn_x)
  
          # prob, mixed_prob = torch.chunk(prob, 2, dim=0)
          prob = torch.chunk(prob, num_domains, dim=0)
          # mixed_prob = torch.chunk(mixed_prob, num_domains, dim=0)
          # x, mixed_x = torch.chunk(x, 2, dim=0)
          x = torch.chunk(x, num_domains, dim=0)
          # mixed_x = torch.chunk(mixed_x, num_domains, dim=0)
  
          return prob, x
  
      def reset_params(self):
          for m in self.modules():
              if isinstance(m, nn.Conv2d):
                  init.kaiming_normal_(m.weight, mode='fan_out')
                  if m.bias is not None:
                      init.constant_(m.bias, 0)
              elif isinstance(m, nn.BatchNorm2d):
                  init.constant_(m.weight, 1)
                  init.constant_(m.bias, 0)
              elif isinstance(m, nn.BatchNorm1d):
                  init.constant_(m.weight, 1)
                  init.constant_(m.bias, 0)
              elif isinstance(m, nn.Linear):
                  init.normal_(m.weight, std=0.001)
                  if m.bias is not None:
                      init.constant_(m.bias, 0)
  
      def get_params(self):
          for param in self.parameters():
              if param.requires_grad:
                  yield param
  
      # def train(self, mode=True):
      #     """
      #     Override the default train() to freeze the BN parameters
      #     """
      #     super().train(mode)
      #     self.freeze_bn()
      #
      # def freeze_bn(self):
      #     for m in self.modules():
      #         if isinstance(m, nn.BatchNorm1d):
      #             m.eval()
      #         if isinstance(m, nn.BatchNorm2d):
      #             m.eval()
  
  
  def resnet18(**kwargs):
      return ResNet(18, **kwargs)
  
  
  def resnet34(**kwargs):
      return ResNet(34, **kwargs)
  
  
  def resnet50(**kwargs):
      return ResNet(50, **kwargs)
  
  
  def resnet101(**kwargs):
      return ResNet(101, **kwargs)
  
  
  def resnet152(**kwargs):
      return ResNet(152, **kwargs)
  
  
  def resnet50_mde(**kwargs):
      return ResNet(50, **kwargs)
  
  
  def resnet50_mldg_smm(**kwargs):
      return ResNet(50, **kwargs)
  
  
apis: torch.nn.init.kaiming_normal_ torch.nn.Linear torch.nn.init.constant_ torch.nn.BatchNorm1d torch.chunk torch.nn.AdaptiveAvgPool2d torch.nn.functional.normalize torch.nn.init.normal_ torch.nn.functional.relu torch.cat torch.nn.Dropout