text: # -*- coding: utf-8 -*-
  
  """
  @date: 2022/4/25 下午6:46
  @file: resnet.py
  @author: zj
  @description: Custom ResNet, derived from torchvision
  """
  from functools import partial
  
  from typing import Type, Any, Callable, Union, List, Optional, Dict, Tuple
  
  import torch
  from torch import nn as nn, Tensor
  from torchvision.models.resnet import BasicBlock, Bottleneck, model_urls
  from torchvision.models.resnet import ResNet as TResNet
  
  # See vision/torchvision/_internally_replaced_utils.py
  # https://github.com/pytorch/vision/blob/b50ffef5f85029b1440ac155ca1e6d95c55520aa/torchvision/_internally_replaced_utils.py
  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
  
  from zcls2.config.key_word import KEY_OUTPUT
  from simpleir.configs.key_words import KEY_FEAT
  
  __all__ = [
      'ResNet',
      'resnet18', 'resnet34', 'resnet50', 'resnet101', 'resnet152',
      'resnext50_32x4d', 'resnext101_32x8d', 'wide_resnet50_2', 'wide_resnet101_2'
  ]
  
  
  class ResNet(TResNet):
      _feat_list = [
          'layer4',
          'avgpool',
          'fc',
      ]
  
      def __init__(self, block: Type[Union[BasicBlock, Bottleneck]], layers: List[int], num_classes: int = 1000,
                   zero_init_residual: bool = False, groups: int = 1, width_per_group: int = 64,
                   replace_stride_with_dilation: Optional[List[bool]] = None,
                   norm_layer: Optional[Callable[..., nn.Module]] = None,
                   feat_type='avgpool') -> None:
          super().__init__(block, layers, num_classes, zero_init_residual, groups, width_per_group,
                           replace_stride_with_dilation, norm_layer)
          assert feat_type in self._feat_list
  
          self.feature_modules = {
              "layer4": self.layer4[2].relu,
              "avgpool": self.avgpool,
              "fc": self.fc
          }
          self.feature_buffer = dict()
          self.feat_type = feat_type
          self._register_hook()
  
      def _register_hook(self) -> None:
          """
          Register hooks to output inner feature map.
          """
  
          def hook(feature_buffer, fea_name, module, input, output):
              feature_buffer[fea_name] = output.data
  
          for fea_name in self._feat_list:
              assert fea_name in self.feature_modules, 'unknown feature {}!'.format(fea_name)
              self.feature_modules[fea_name].register_forward_hook(partial(hook, self.feature_buffer, fea_name))
  
      def forward(self, x: Tensor) -> Dict:
          res = super(ResNet, self).forward(x)
          return {
              KEY_OUTPUT: res,
              KEY_FEAT: self.feature_buffer[self.feat_type]
          }
  
  
  def _resnet(
          arch: str,
          block: Type[Union[BasicBlock, Bottleneck]],
          layers: List[int],
          pretrained: bool,
          progress: bool,
          **kwargs: Any
  ) -> ResNet:
      model = ResNet(block, layers, **kwargs)
  
      if pretrained:
          state_dict = load_state_dict_from_url(model_urls[arch],
                                                progress=progress)
  
          # If the number of model outputs is different from the model setting,
          # the corresponding pretraining model weight will not be loaded
          if model.fc.out_features != 1000:
              state_dict.pop('fc.weight')
              state_dict.pop('fc.bias')
  
          ret = model.load_state_dict(state_dict, strict=False)
          assert set(ret.missing_keys) == {'fc.weight', 'fc.bias'}, \
              f'Missing keys when loading pretrained weights: {ret.missing_keys}'
  
      return model
  
  
  def resnet18(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:
      r"""ResNet-18 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('resnet18', BasicBlock, [2, 2, 2, 2], pretrained, progress,
                     **kwargs)
  
  
  def resnet34(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:
      r"""ResNet-34 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('resnet34', BasicBlock, [3, 4, 6, 3], pretrained, progress,
                     **kwargs)
  
  
  def resnet50(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:
      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 resnet101(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:
      r"""ResNet-101 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('resnet101', Bottleneck, [3, 4, 23, 3], pretrained, progress,
                     **kwargs)
  
  
  def resnet152(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:
      r"""ResNet-152 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('resnet152', Bottleneck, [3, 8, 36, 3], pretrained, progress,
                     **kwargs)
  
  
  def resnext50_32x4d(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:
      r"""ResNeXt-50 32x4d model from
      `"Aggregated Residual Transformation for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.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
      """
      kwargs['groups'] = 32
      kwargs['width_per_group'] = 4
      return _resnet('resnext50_32x4d', Bottleneck, [3, 4, 6, 3],
                     pretrained, progress, **kwargs)
  
  
  def resnext101_32x8d(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:
      r"""ResNeXt-101 32x8d model from
      `"Aggregated Residual Transformation for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.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
      """
      kwargs['groups'] = 32
      kwargs['width_per_group'] = 8
      return _resnet('resnext101_32x8d', Bottleneck, [3, 4, 23, 3],
                     pretrained, progress, **kwargs)
  
  
  def wide_resnet50_2(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:
      r"""Wide ResNet-50-2 model from
      `"Wide Residual Networks" <https://arxiv.org/pdf/1605.07146.pdf>`_.
  
      The model is the same as ResNet except for the bottleneck number of channels
      which is twice larger in every block. The number of channels in outer 1x1
      convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048
      channels, and in Wide ResNet-50-2 has 2048-1024-2048.
  
      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
      """
      kwargs['width_per_group'] = 64 * 2
      return _resnet('wide_resnet50_2', Bottleneck, [3, 4, 6, 3],
                     pretrained, progress, **kwargs)
  
  
  def wide_resnet101_2(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNet:
      r"""Wide ResNet-101-2 model from
      `"Wide Residual Networks" <https://arxiv.org/pdf/1605.07146.pdf>`_.
  
      The model is the same as ResNet except for the bottleneck number of channels
      which is twice larger in every block. The number of channels in outer 1x1
      convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048
      channels, and in Wide ResNet-50-2 has 2048-1024-2048.
  
      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
      """
      kwargs['width_per_group'] = 64 * 2
      return _resnet('wide_resnet101_2', Bottleneck, [3, 4, 23, 3],
                     pretrained, progress, **kwargs)
  