text: 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
  
  
  __all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
             'resnet152', 'resnet50_mde']
  
  
  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
          self.num_domains = len(num_classes)
  
          # 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)
  
          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.classifiers = nn.ModuleList()
              for domain_id, num_classes in enumerate(self.num_classes):
                  classifier = nn.Linear(self.num_features, num_classes, bias=False)
                  init.normal_(classifier.weight, std=0.001)
                  self.classifiers.append(classifier)
  
          init.constant_(self.feat_bn.weight, 1)
          init.constant_(self.feat_bn.bias, 0)
  
          if not pretrained:
              self.reset_params()
  
      def forward(self, x: list, output_prob=False):
          if self.training:
              x = torch.cat(x, dim=0)
  
          x = self.conv(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)
  
          probs = []
          all_x = torch.chunk(bn_x, self.num_domains, dim=0)
          for domain_id, in_x in enumerate(all_x):
              prob = self.classifiers[domain_id](in_x)
              probs.append(prob)
  
          x = torch.chunk(x, self.num_domains, dim=0)
  
          return probs, 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 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)
  
  