text: import torch
  import torch.nn as nn
  import pdb
  import os
  import torch.nn.functional as F
  import math
  
  def initialize_weights(modules, init_mode):
      for m in modules:
          if isinstance(m, nn.Conv2d):
              if init_mode == 'he':
                  nn.init.kaiming_normal_(m.weight, mode='fan_out',
                                          nonlinearity='relu')
              elif init_mode == 'xavier':
                  nn.init.xavier_uniform_(m.weight.data)
              else:
                  raise ValueError('Invalid init_mode {}'.format(init_mode))
              if m.bias is not None:
                  nn.init.constant_(m.bias, 0)
          elif isinstance(m, nn.BatchNorm2d):
              nn.init.constant_(m.weight, 1)
              nn.init.constant_(m.bias, 0)
          elif isinstance(m, nn.Linear):
              nn.init.normal_(m.weight, 0, 0.01)
              nn.init.constant_(m.bias, 0)
  
  
  def remove_layer(state_dict, keyword):
      keys = [key for key in state_dict.keys()]
      for key in keys:
          if keyword in key:
              state_dict.pop(key)
      return state_dict
  
  class resnet_layer4(nn.Module):
      def __init__(self):
          super().__init__()
  
          self.features_7_0_conv1 = nn.Conv2d(1024, 512, kernel_size=[1, 1], stride=(1, 1), bias=False)
          self.features_7_0_bn1 = nn.BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True)
          self.features_7_0_relu1 = nn.ReLU()
          self.features_7_0_conv2 = nn.Conv2d(512, 512, kernel_size=[3, 3], stride=(2, 2), padding=(1, 1), bias=False)
          self.features_7_0_bn2 = nn.BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True)
          self.features_7_0_relu2 = nn.ReLU()
          self.features_7_0_conv3 = nn.Conv2d(512, 2048, kernel_size=[1, 1], stride=(1, 1), bias=False)
          self.features_7_0_bn3 = nn.BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True)
          self.features_7_0_downsample_0 = nn.Conv2d(1024, 2048, kernel_size=[1, 1], stride=(2, 2), bias=False)
          self.features_7_0_downsample_1 = nn.BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True)
          self.features_7_0_id_relu = nn.ReLU()
          self.features_7_1_conv1 = nn.Conv2d(2048, 512, kernel_size=[1, 1], stride=(1, 1), bias=False)
          self.features_7_1_bn1 = nn.BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True)
          self.features_7_1_relu1 = nn.ReLU()
          self.features_7_1_conv2 = nn.Conv2d(512, 512, kernel_size=[3, 3], stride=(1, 1), padding=(1, 1), bias=False)
          self.features_7_1_bn2 = nn.BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True)
          self.features_7_1_relu2 = nn.ReLU()
          self.features_7_1_conv3 = nn.Conv2d(512, 2048, kernel_size=[1, 1], stride=(1, 1), bias=False)
          self.features_7_1_bn3 = nn.BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True)
          self.features_7_1_id_relu = nn.ReLU()
          self.features_7_2_conv1 = nn.Conv2d(2048, 512, kernel_size=[1, 1], stride=(1, 1), bias=False)
          self.features_7_2_bn1 = nn.BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True)
          self.features_7_2_relu1 = nn.ReLU()
          self.features_7_2_conv2 = nn.Conv2d(512, 512, kernel_size=[3, 3], stride=(1, 1), padding=(1, 1), bias=False)
          self.features_7_2_bn2 = nn.BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True)
          self.features_7_2_relu2 = nn.ReLU()
          self.features_7_2_conv3 = nn.Conv2d(512, 2048, kernel_size=[1, 1], stride=(1, 1), bias=False)
          self.features_7_2_bn3 = nn.BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True)
          self.features_7_2_id_relu = nn.ReLU()
          self.features_8 = nn.AvgPool2d(kernel_size=[7, 7], stride=[1, 1], padding=0, ceil_mode=False, count_include_pad=False)
  
      def forward(self, data):
          
          features_7_0_conv1 = self.features_7_0_conv1(data)
          features_7_0_bn1 = self.features_7_0_bn1(features_7_0_conv1)
          features_7_0_relu1 = self.features_7_0_relu1(features_7_0_bn1)
          features_7_0_conv2 = self.features_7_0_conv2(features_7_0_relu1)
          features_7_0_bn2 = self.features_7_0_bn2(features_7_0_conv2)
          features_7_0_relu2 = self.features_7_0_relu2(features_7_0_bn2)
          features_7_0_conv3 = self.features_7_0_conv3(features_7_0_relu2)
          features_7_0_bn3 = self.features_7_0_bn3(features_7_0_conv3)
          features_7_0_downsample_0 = self.features_7_0_downsample_0(data)
          features_7_0_downsample_1 = self.features_7_0_downsample_1(features_7_0_downsample_0)
          features_7_0_merge = torch.add(features_7_0_downsample_1, 1, features_7_0_bn3)
          features_7_0_id_relu = self.features_7_0_id_relu(features_7_0_merge)
          features_7_1_conv1 = self.features_7_1_conv1(features_7_0_id_relu)
          features_7_1_bn1 = self.features_7_1_bn1(features_7_1_conv1)
          features_7_1_relu1 = self.features_7_1_relu1(features_7_1_bn1)
          features_7_1_conv2 = self.features_7_1_conv2(features_7_1_relu1)
          features_7_1_bn2 = self.features_7_1_bn2(features_7_1_conv2)
          features_7_1_relu2 = self.features_7_1_relu2(features_7_1_bn2)
          features_7_1_conv3 = self.features_7_1_conv3(features_7_1_relu2)
          features_7_1_bn3 = self.features_7_1_bn3(features_7_1_conv3)
          features_7_1_merge = torch.add(features_7_0_id_relu, 1, features_7_1_bn3)
          features_7_1_id_relu = self.features_7_1_id_relu(features_7_1_merge)
          features_7_2_conv1 = self.features_7_2_conv1(features_7_1_id_relu)
          features_7_2_bn1 = self.features_7_2_bn1(features_7_2_conv1)
          features_7_2_relu1 = self.features_7_2_relu1(features_7_2_bn1)
          features_7_2_conv2 = self.features_7_2_conv2(features_7_2_relu1)
          features_7_2_bn2 = self.features_7_2_bn2(features_7_2_conv2)
          features_7_2_relu2 = self.features_7_2_relu2(features_7_2_bn2)
          features_7_2_conv3 = self.features_7_2_conv3(features_7_2_relu2)
          features_7_2_bn3 = self.features_7_2_bn3(features_7_2_conv3)
          features_7_2_merge = torch.add(features_7_1_id_relu, 1, features_7_2_bn3)
          features_7_2_id_relu = self.features_7_2_id_relu(features_7_2_merge)
          features_8 = self.features_8(features_7_2_id_relu)
          classifier_flatten = features_8.view(features_8.size(0), -1)
          # classifier_0 = self.classifier_0(classifier_flatten)
          return classifier_flatten, features_7_2_id_relu      
  
  
  
  
  
  
  
  
  
  
  
  
  
  
  # import torch
  # import torch.nn as nn
  # # from .utils import load_state_dict_from_url
  
  
  # __all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
  #            'resnet152', 'resnext50_32x4d', 'resnext101_32x8d',
  #            'wide_resnet50_2', 'wide_resnet101_2']
  
  
  # model_urls = {
  #     'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
  #     'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
  #     'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
  #     'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
  #     'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
  #     'resnext50_32x4d': 'https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth',
  #     'resnext101_32x8d': 'https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pth',
  #     'wide_resnet50_2': 'https://download.pytorch.org/models/wide_resnet50_2-95faca4d.pth',
  #     'wide_resnet101_2': 'https://download.pytorch.org/models/wide_resnet101_2-32ee1156.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=dilation, groups=groups, bias=False, dilation=dilation)
  
  
  # def conv1x1(in_planes, out_planes, stride=1):
  #     """1x1 convolution"""
  #     return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
  
  
  # class BasicBlock(nn.Module):
  #     expansion = 1
  
  #     def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
  #                  base_width=64, dilation=1, norm_layer=None):
  #         super(BasicBlock, self).__init__()
  #         if norm_layer is None:
  #             norm_layer = nn.BatchNorm2d
  #         if groups != 1 or base_width != 64:
  #             raise ValueError('BasicBlock only supports groups=1 and base_width=64')
  #         if dilation > 1:
  #             raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
  #         # Both self.conv1 and self.downsample layers downsample the input when stride != 1
  #         self.conv1 = conv3x3(inplanes, planes, stride)
  #         self.bn1 = norm_layer(planes)
  #         self.relu = nn.ReLU(inplace=True)
  #         self.conv2 = conv3x3(planes, planes)
  #         self.bn2 = norm_layer(planes)
  #         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)
  
  #         if self.downsample is not None:
  #             identity = self.downsample(x)
  
  #         out += identity
  #         out = self.relu(out)
  
  #         return out
  
  
  # 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 Bottleneck(nn.Module):
      expansion = 4
  
      def __init__(self, inplanes, planes, stride=1, downsample=None,
                   base_width=64):
          super(Bottleneck, self).__init__()
          width = int(planes * (base_width / 64.))
          self.conv1 = nn.Conv2d(inplanes, width, 1, bias=False)
          self.bn1 = nn.BatchNorm2d(width)
          self.conv2 = nn.Conv2d(width, width, 3,
                                 stride=stride, padding=1, bias=False)
          self.bn2 = nn.BatchNorm2d(width)
          self.conv3 = nn.Conv2d(width, planes * self.expansion, 1, bias=False)
          self.bn3 = nn.BatchNorm2d(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], stride=1 )
  #         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=1,
  #                                        dilate=replace_stride_with_dilation[1])
  #         self.layer4 = self._make_layer(block, 512, layers[3], stride=1,
  #                                        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)
  
  
  
  
  
  
  class ResNetCam(nn.Module):
      def __init__(self, block, layers, num_classes=1000,
                   large_feature_map=False, **kwargs):
          super(ResNetCam, self).__init__()
  
          stride_l3 = 1 if large_feature_map else 2
          self.inplanes = 64
  
          self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2,
                                 padding=3, bias=False)
          self.bn1 = nn.BatchNorm2d(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], stride=1)
          self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
          self.layer3 = self._make_layer(block, 256, layers[2], stride=stride_l3)
          self.layer4 = self._make_layer(block, 512, layers[3], stride=1)
          self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
          self.fc = nn.Linear(512 * block.expansion, num_classes)
  
          initialize_weights(self.modules(), init_mode='xavier')
  
      def forward(self, x, labels=None, return_cam=False):
          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)
  
          # pre_logit = self.avgpool(x)
          # pre_logit = pre_logit.reshape(pre_logit.size(0), -1)
          # logits = self.fc(pre_logit)
  
          # if return_cam:
          #     feature_map = x.detach().clone()
          #     cam_weights = self.fc.weight[labels]
          #     cams = (cam_weights.view(*feature_map.shape[:2], 1, 1) *
          #             feature_map).mean(1, keepdim=False)
          #     return cams
          # return {'logits': logits}
  
          return x
  
      def _make_layer(self, block, planes, blocks, stride):
          layers = self._layer(block, planes, blocks, stride)
          return nn.Sequential(*layers)
  
      def _layer(self, block, planes, blocks, stride):
          downsample = get_downsampling_layer(self.inplanes, block, planes,
                                              stride)
  
          layers = [block(self.inplanes, planes, stride, downsample)]
          self.inplanes = planes * block.expansion
          for _ in range(1, blocks):
              layers.append(block(self.inplanes, planes))
  
          return layers
  
  
  
  class ResNetCam_I2C(nn.Module):
  
      def __init__(self, block, layers, num_classes=1000, args=None):
          self.inplanes = 64
          super(ResNetCam_I2C, self).__init__()
          self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
                                 bias=False)
          self.bn1 = nn.BatchNorm2d(64)
          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)
          self.layer3 = self._make_layer(block, 256, layers[2], stride=1)
          self.layer4 = self._make_layer(block, 512, layers[3], stride=1)
  
  
          #Added
  
          # pdb.set_trace()
  
  
          # self.loss_local_factor = args.loss_local_factor
          # self.local_seed_num = args.local_seed_num
          # self.loss_global_factor = args.loss_global_factor
          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_()
  
      def _make_layer(self, block, planes, blocks, stride=1):
          downsample = None
          if stride != 1 or self.inplanes != planes * block.expansion:
              downsample = nn.Sequential(
                  nn.Conv2d(self.inplanes, planes * block.expansion,
                            kernel_size=1, stride=stride, bias=False),
                  nn.BatchNorm2d(planes * block.expansion),
              )
  
          layers = []
          layers.append(block(self.inplanes, planes, stride, downsample))
          self.inplanes = planes * block.expansion
          for i in range(1, blocks):
              layers.append(block(self.inplanes, planes))
  
          return nn.Sequential(*layers)
  
      def forward(self, x, gt_labels=None):
          x = self.conv1(x)
          x = self.bn1(x)
          x = self.relu(x)
          x = self.maxpool(x)
  
          feat1 = self.layer1(x)
          feat2 = self.layer2(feat1)
          feat3 = self.layer3(feat2)
          feat4 = self.layer4(feat3)
  
          #Branch 1
  
          return feat4
  
  
  
  
  
  
  
  
  
  
  
  class ResNetCam_ori(nn.Module):
      def __init__(self, block, layers, num_classes=1000,
                   large_feature_map=False, **kwargs):
          super(ResNetCam_ori, self).__init__()
  
          stride_l3 = 1 if large_feature_map else 2
          self.inplanes = 64
  
          self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2,
                                 padding=3, bias=False)
          self.bn1 = nn.BatchNorm2d(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], stride=1)
          self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
          self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
          self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
          self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
          self.fc = nn.Linear(512 * block.expansion, num_classes)
  
          initialize_weights(self.modules(), init_mode='xavier')
  
      def forward(self, x, labels=None, return_cam=False):
          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)
  
          # pre_logit = self.avgpool(x)
          # pre_logit = pre_logit.reshape(pre_logit.size(0), -1)
          # logits = self.fc(pre_logit)
  
          # if return_cam:
          #     feature_map = x.detach().clone()
          #     cam_weights = self.fc.weight[labels]
          #     cams = (cam_weights.view(*feature_map.shape[:2], 1, 1) *
          #             feature_map).mean(1, keepdim=False)
          #     return cams
          # return {'logits': logits}
  
          return x
  
      def _make_layer(self, block, planes, blocks, stride):
          layers = self._layer(block, planes, blocks, stride)
          return nn.Sequential(*layers)
  
      def _layer(self, block, planes, blocks, stride):
          downsample = get_downsampling_layer(self.inplanes, block, planes,
                                              stride)
  
          layers = [block(self.inplanes, planes, stride, downsample)]
          self.inplanes = planes * block.expansion
          for _ in range(1, blocks):
              layers.append(block(self.inplanes, planes))
  
          return layers
  
  
  
  
  
  def _resnet(arch, block, layers, pretrained, progress, **kwargs):
      model = ResNet(block, layers, **kwargs)
      if pretrained:
          state_dict = '/esat/monkey/kwang/pytorchModel/resnet50-19c8e357.pth'
  
          # remove_layer(state_dict, 'Mixed_7')
          # remove_layer(state_dict, 'AuxLogits')
          # remove_layer(state_dict, 'fc.')
  
  
          model.load_state_dict(torch.load(state_dict))
      return model
  
  
  
  
  def resnet50(pretrained=False, progress=True, **kwargs):
      r"""ResNet-50 model from
      `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
      Args:
          pretrained (bool): If True, returns a model pre-trained on ImageNet
          progress (bool): If True, displays a progress bar of the download to stderr
      """
      return _resnet('resnet50', Bottleneck, [3, 4, 6, 3], pretrained, progress,
                     **kwargs)
  
  def get_downsampling_layer(inplanes, block, planes, stride):
      outplanes = planes * block.expansion
      if stride == 1 and inplanes == outplanes:
          return
      else:
          return nn.Sequential(
              nn.Conv2d(inplanes, outplanes, 1, stride, bias=False),
              nn.BatchNorm2d(outplanes),
          )
  
  ##############################################################################################
  
  def resnet50_cvpr(architecture_type, pretrained=False, pretrained_path=None, ori=False,
               **kwargs):
  
      if ori is False:
          model = {'cam': ResNetCam}[architecture_type](Bottleneck, [3, 4, 6, 3], large_feature_map=True,
                                                    **kwargs)
      else:
  
          model = {'cam': ResNetCam_ori}[architecture_type](Bottleneck, [3, 4, 6, 3], large_feature_map=True,
                                                    **kwargs)
  
  
  
      if pretrained:
          model = load_pretrained_model_resnet(model, architecture_type,
                                        path=pretrained_path, **kwargs)
      return model
  
  
  
  def resnet50_i2c(numberclass = 1000, pretrained=False,  **kwargs):
      """Constructs a ResNet-50 model.
  
      Args:
          pretrained (bool): If True, returns a model pre-trained on ImageNet
      """
      state_dict = torch.load('/esat/monkey/kwang/pytorchModel/resnet50-19c8e357.pth')
  
      model = ResNetCam_I2C(Bottleneck, [3, 4, 6, 3], numberclass, **kwargs)
      if pretrained:
          model.load_state_dict(state_dict, strict=False)
      return model
  
  
  def load_pretrained_model_resnet(model, wsol_method, path=None, **kwargs):
      strict_rule = True
  
      state_dict = torch.load('/esat/monkey/kwang/pytorchModel/resnet50-19c8e357.pth')
   
  
  
      # if kwargs['dataset_name'] != 'ILSVRC' or wsol_method in ('acol', 'spg'):
      state_dict = remove_layer(state_dict, 'fc')
      strict_rule = False
  
      model.load_state_dict(state_dict, strict=strict_rule)
      return model
  
  
  
  
  
  
  
  
  
  
  
  
  
  
  
  class CLUB(nn.Module):  # CLUB: Mutual Information Contrastive Learning Upper Bound
      '''
          This class provides the CLUB estimation to I(X,Y)
          Method:
              mi_est() :      provides the estimation with input samples  
              loglikeli() :   provides the log-likelihood of the approximation q(Y|X) with input samples
          Arguments:
              x_dim, y_dim :         the dimensions of samples from X, Y respectively
              hidden_size :          the dimension of the hidden layer of the approximation network q(Y|X)
              x_samples, y_samples : samples from X and Y, having shape [sample_size, x_dim/y_dim] 
      '''
      def __init__(self, x_dim, y_dim):
          super(CLUB, self).__init__()
          # p_mu outputs mean of q(Y|X)
          self.p_mu = nn.Sequential(nn.Linear(x_dim,x_dim // 2),
                                         nn.ReLU(),
                                         nn.Linear(x_dim,x_dim),
                                         nn.ReLU(),
                                         nn.Linear(x_dim // 2, y_dim))
          # p_logvar outputs log of variance of q(Y|X)
          self.p_logvar = nn.Sequential(nn.Linear(x_dim, x_dim//2),
                                          nn.ReLU(),
                                          nn.Linear(x_dim, x_dim),
                                          nn.ReLU(),
                                          nn.Linear(x_dim//2, y_dim),
                                          nn.Tanh()
                                          )
  
      def get_mu_logvar(self, x_samples):
          mu = self.p_mu(x_samples)
          logvar = self.p_logvar(x_samples)
          return mu, logvar
      
      def mi_est(self, x_samples, y_samples): 
          mu, logvar = self.get_mu_logvar(x_samples)
          
          # log of conditional probability of positive sample pairs
          positive = - (mu - y_samples)**2 /2./logvar.exp()  
          
          prediction_1 = mu.unsqueeze(1)          # shape [nsample,1,dim]
          y_samples_1 = y_samples.unsqueeze(0)    # shape [1,nsample,dim]
  
          # log of conditional probability of negative sample pairs
          # pdb.set_trace()
          negative = - ((y_samples_1 - prediction_1)**2).mean(dim=1)/2./logvar.exp() 
  
          return (positive.sum(dim = -1) - negative.sum(dim = -1)).mean()
  
      def loglikeli(self, x_samples, y_samples): # unnormalized loglikelihood 
          mu, logvar = self.get_mu_logvar(x_samples)
          return (-(mu - y_samples)**2 /logvar.exp()-logvar).sum(dim=1).mean(dim=0)
  
  
  
  
  
  
  
  
  
  
  
  def load_pretrained_model(model, path=None):
      state_dict = torch.load(path)
  
  
      # remove_layer(state_dict, 'Mixed_7')
      remove_layer(state_dict, 'AuxLogits')
      remove_layer(state_dict, 'fc.')
      # remove_layer(state_dict, 'fc.')
  
      model.load_state_dict(state_dict, strict=False)
      print('Load Pretrained model')
      return model
  
  
  def inception_v3(architecture_type, pretrained=False, pretrained_path=None,
                   **kwargs):
      model = {'cam': InceptionCam
               }[architecture_type](**kwargs)
      # if pretrained:
      # model = load_pretrained_model(model, '/esat/monkey/kwang/pytorchModel/inception_v3_google-1a9a5a14.pth')
      model = load_pretrained_model(model, '/esat/monkey/kwang/pytorchModel/inception_v3_google-1a9a5a14.pth')
      print('loaded')
  
      return model
  
  
  
  
  def _make_divisible(v, divisor, min_value=None):
      """
      This function is taken from the original tf repo.
      It ensures that all layers have a channel number that is divisible by 8
      It can be seen here:
      https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py
      :param v:
      :param divisor:
      :param min_value:
      :return:
      """
      if min_value is None:
          min_value = divisor
      new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
      # Make sure that round down does not go down by more than 10%.
      if new_v < 0.9 * v:
          new_v += divisor
      return new_v
  
  
  class ConvBNReLU(nn.Sequential):
      def __init__(self, in_planes, out_planes, kernel_size=3, stride=1, groups=1, norm_layer=None):
          padding = (kernel_size - 1) // 2
          if norm_layer is None:
              norm_layer = nn.BatchNorm2d
          super(ConvBNReLU, self).__init__(
              nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, groups=groups, bias=False),
              norm_layer(out_planes),
              nn.ReLU6(inplace=True)
          )
  
  
  class InvertedResidual(nn.Module):
      def __init__(self, inp, oup, stride, counter, expand_ratio, norm_layer=None):
          super(InvertedResidual, self).__init__()
          self.stride = stride
          assert stride in [1, 2]
  
          if norm_layer is None:
              norm_layer = nn.BatchNorm2d
  
          hidden_dim = int(round(inp * expand_ratio))
          self.use_res_connect = self.stride == 1 and inp == oup
          # print(self.use_res_connect)
          # pdb.set_trace()
          # if counter > 13:
          #     self.use_res_connect = False
          layers = []
          if expand_ratio != 1:
              # pw
              layers.append(ConvBNReLU(inp, hidden_dim, kernel_size=1, norm_layer=norm_layer))
          layers.extend([
              # dw
              ConvBNReLU(hidden_dim, hidden_dim, stride=stride, groups=hidden_dim, norm_layer=norm_layer),
              # pw-linear
              nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
              norm_layer(oup),
          ])
          self.conv = nn.Sequential(*layers)
  
      def forward(self, x):
          if self.use_res_connect:
              return x + self.conv(x)
          else:
              return self.conv(x)
  
      # def forward(self, x):
      #     if self.use_res_connect:
      #         return self.conv(x)
      #     else:
      #         return self.conv(x)
  
  
  class MobileNetV2(nn.Module):
      def __init__(self,
                   num_classes=1000,
                   width_mult=1.0,
                   inverted_residual_setting=None,
                   round_nearest=8,
                   block=None,
                   norm_layer=None):
          """
          MobileNet V2 main class
          Args:
              num_classes (int): Number of classes
              width_mult (float): Width multiplier - adjusts number of channels in each layer by this amount
              inverted_residual_setting: Network structure
              round_nearest (int): Round the number of channels in each layer to be a multiple of this number
              Set to 1 to turn off rounding
              block: Module specifying inverted residual building block for mobilenet
              norm_layer: Module specifying the normalization layer to use
          """
          super(MobileNetV2, self).__init__()
  
          if block is None:
              block = InvertedResidual
  
          if norm_layer is None:
              norm_layer = nn.BatchNorm2d
  
          input_channel = 32
          last_channel = 1280
  
          if inverted_residual_setting is None:
              inverted_residual_setting = [
                  # t, c, n, s
                  [1, 16, 1, 1],
                  [6, 24, 2, 2],
                  [6, 32, 3, 2],
                  [6, 64, 4, 1],
                  [6, 96, 3, 1],
                  [6, 160, 3, 1],
                  [6, 320, 1, 1],
              ]
  
          # only check the first element, assuming user knows t,c,n,s are required
          if len(inverted_residual_setting) == 0 or len(inverted_residual_setting[0]) != 4:
              raise ValueError("inverted_residual_setting should be non-empty "
                               "or a 4-element list, got {}".format(inverted_residual_setting))
  
          # building first layer
          input_channel = _make_divisible(input_channel * width_mult, round_nearest)
          self.last_channel = _make_divisible(last_channel * max(1.0, width_mult), round_nearest)
          features = [ConvBNReLU(3, input_channel, stride=2, norm_layer=norm_layer)]
          self.counter = 0
          # building inverted residual blocks
          for t, c, n, s in inverted_residual_setting:
              output_channel = _make_divisible(c * width_mult, round_nearest)
              for i in range(n):
                  stride = s if i == 0 else 1
                  features.append(block(input_channel, output_channel, stride, self.counter, expand_ratio=t, norm_layer=norm_layer))
                  self.counter += 1
                  # print(self.counter)
                  input_channel = output_channel
          # building last several layers
          features.append(ConvBNReLU(input_channel, self.last_channel, kernel_size=1, norm_layer=norm_layer))
          # make it nn.Sequential
          self.features = nn.Sequential(*features)
  
          # building classifier
          self.classifier = nn.Sequential(
              nn.Dropout(0.2),
              nn.Linear(self.last_channel, num_classes),
          )
  
          # weight initialization
          for m in self.modules():
              if isinstance(m, nn.Conv2d):
                  nn.init.kaiming_normal_(m.weight, mode='fan_out')
                  if m.bias is not None:
                      nn.init.zeros_(m.bias)
              elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
                  nn.init.ones_(m.weight)
                  nn.init.zeros_(m.bias)
              elif isinstance(m, nn.Linear):
                  nn.init.normal_(m.weight, 0, 0.01)
                  nn.init.zeros_(m.bias)
  
      def _forward_impl(self, x):
          # This exists since TorchScript doesn't support inheritance, so the superclass method
          # (this one) needs to have a name other than `forward` that can be accessed in a subclass
          # pdb.set_trace()
          x = self.features(x)
          # Cannot use "squeeze" as batch-size can be 1 => must use reshape with x.shape[0]
          # pdb.set_trace()
          # x = nn.functional.adaptive_avg_pool2d(x, 1).reshape(x.shape[0], -1)
          # x = self.classifier(x)
          return x
  
      def forward(self, x):
          # pdb.set_trace()
          return self._forward_impl(x)
  
  
  class MobileNetV2_ori(nn.Module):
      def __init__(self,
                   num_classes=1000,
                   width_mult=1.0,
                   inverted_residual_setting=None,
                   round_nearest=8,
                   block=None,
                   norm_layer=None):
          """
          MobileNet V2 main class
          Args:
              num_classes (int): Number of classes
              width_mult (float): Width multiplier - adjusts number of channels in each layer by this amount
              inverted_residual_setting: Network structure
              round_nearest (int): Round the number of channels in each layer to be a multiple of this number
              Set to 1 to turn off rounding
              block: Module specifying inverted residual building block for mobilenet
              norm_layer: Module specifying the normalization layer to use
          """
          super(MobileNetV2_ori, self).__init__()
  
          if block is None:
              block = InvertedResidual
  
          if norm_layer is None:
              norm_layer = nn.BatchNorm2d
  
          input_channel = 32
          last_channel = 1280
  
          if inverted_residual_setting is None:
              inverted_residual_setting = [
                  # t, c, n, s
                  [1, 16, 1, 1],
                  [6, 24, 2, 2],
                  [6, 32, 3, 2],
                  [6, 64, 4, 2],
                  [6, 96, 3, 1],
                  [6, 160, 3, 2],
                  [6, 320, 1, 1],
              ]
  
          # only check the first element, assuming user knows t,c,n,s are required
          if len(inverted_residual_setting) == 0 or len(inverted_residual_setting[0]) != 4:
              raise ValueError("inverted_residual_setting should be non-empty "
                               "or a 4-element list, got {}".format(inverted_residual_setting))
  
          # building first layer
          input_channel = _make_divisible(input_channel * width_mult, round_nearest)
          self.last_channel = _make_divisible(last_channel * max(1.0, width_mult), round_nearest)
          features = [ConvBNReLU(3, input_channel, stride=2, norm_layer=norm_layer)]
          # building inverted residual blocks
          for t, c, n, s in inverted_residual_setting:
              output_channel = _make_divisible(c * width_mult, round_nearest)
              for i in range(n):
                  stride = s if i == 0 else 1
                  features.append(block(input_channel, output_channel, stride, 0, expand_ratio=t, norm_layer=norm_layer))
                  input_channel = output_channel
          # building last several layers
          features.append(ConvBNReLU(input_channel, self.last_channel, kernel_size=1, norm_layer=norm_layer))
          # make it nn.Sequential
          self.features = nn.Sequential(*features)
  
          # building classifier
          self.classifier = nn.Sequential(
              nn.Dropout(0.2),
              nn.Linear(self.last_channel, num_classes),
          )
  
          # weight initialization
          for m in self.modules():
              if isinstance(m, nn.Conv2d):
                  nn.init.kaiming_normal_(m.weight, mode='fan_out')
                  if m.bias is not None:
                      nn.init.zeros_(m.bias)
              elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
                  nn.init.ones_(m.weight)
                  nn.init.zeros_(m.bias)
              elif isinstance(m, nn.Linear):
                  nn.init.normal_(m.weight, 0, 0.01)
                  nn.init.zeros_(m.bias)
  
      def _forward_impl(self, x):
          # This exists since TorchScript doesn't support inheritance, so the superclass method
          # (this one) needs to have a name other than `forward` that can be accessed in a subclass
          x = self.features(x)
          # Cannot use "squeeze" as batch-size can be 1 => must use reshape with x.shape[0]
          # pdb.set_trace()
          # x = nn.functional.adaptive_avg_pool2d(x, 1).reshape(x.shape[0], -1)
          # x = self.classifier(x)
          return x
  
      def forward(self, x):
          return self._forward_impl(x)
  
  
  def mobilenet_v2(pretrained=False, progress=True, ori=False, **kwargs):
      """
      Constructs a MobileNetV2 architecture from
      `"MobileNetV2: Inverted Residuals and Linear Bottlenecks" <https://arxiv.org/abs/1801.04381>`_.
      Args:
          pretrained (bool): If True, returns a model pre-trained on ImageNet
          progress (bool): If True, displays a progress bar of the download to stderr
      """
      if ori is False:
          model = MobileNetV2(**kwargs)
          # pdb.set_trace()
      else:
          model = MobileNetV2_ori(**kwargs)
      # if pretrained:
      #     state_dict = load_state_dict_from_url(model_urls['mobilenet_v2'],
      #                                           progress=progress)
      #     model.load_state_dict(state_dict)
  
      # state_dict = load_state_dict()
  
      model.load_state_dict(torch.load('/esat/monkey/kwang/pytorchModel/mobilenet_v2-b0353104.pth'))
      return model
  
  
  
  
  class Net(nn.Module):
      def __init__(self):
          super(Net, self).__init__()
  
          def conv_bn(inp, oup, stride):
              return nn.Sequential(
                  nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
                  nn.BatchNorm2d(oup),
                  nn.ReLU(inplace=True)
              )
  
          def conv_dw(inp, oup, stride):
              return nn.Sequential(
                  nn.Conv2d(inp, inp, 3, stride, 1, groups=inp, bias=False),
                  nn.BatchNorm2d(inp),
                  nn.ReLU(inplace=True),
      
                  nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
                  nn.BatchNorm2d(oup),
                  nn.ReLU(inplace=True),
              )
  
  
  
          self.model = nn.Sequential(
              conv_bn(  3,  32, 2), 
              conv_dw( 32,  64, 1),
              conv_dw( 64, 128, 2),
              conv_dw(128, 128, 1),
              conv_dw(128, 256, 2),
              conv_dw(256, 256, 1),
              conv_dw(256, 512, 1),
              conv_dw(512, 512, 1),
              conv_dw(512, 512, 1),
              conv_dw(512, 512, 1),
              conv_dw(512, 512, 1),
              conv_dw(512, 512, 1),
              conv_dw(512, 1024, 1),
              conv_dw(1024, 1024, 1),
              # nn.AvgPool2d(7),
          )
          self.fc = nn.Linear(1024, 1000)
  
          # self.module = self.model
  
      def forward(self, x):
          x = self.model(x)
          # x = x.view(-1, 1024)
          # x = self.fc(x)
          return x
  
  
  
  def mobilenet_v1(path='/users/visics/kwang/Desktop/Works/WSOD/pytorch-mobilenet-v1/mobilenet_sgd_68.848.pth.tar'):
      net = Net()
      # pdb.set_trace()
      state_dict = torch.load(path)
      useful_dict = []
  
      aaa = [*state_dict['state_dict'].keys()]
  
      for i in range(len(state_dict['state_dict'].keys())):
          state_dict['state_dict']['{}'.format(aaa[i][7:])] = state_dict['state_dict']['{}'.format(aaa[i])]
          del state_dict['state_dict']['{}'.format(aaa[i])]
  
  
      # pdb.set_trace()
  
      net.load_state_dict(state_dict['state_dict'])
      return net