text: # Copyright (c) Microsoft Corporation. All rights reserved.
  # Licensed under the MIT License.
  """
  Variant of the resnet module that takes cfg as an argument.
  Example usage. Strings may be specified in the config file.
      model = ResNet(
          "StemWithFixedBatchNorm",
          "BottleneckWithFixedBatchNorm",
          "ResNet50StagesTo4",
      )
  OR:
      model = ResNet(
          "StemWithGN",
          "BottleneckWithGN",
          "ResNet50StagesTo4",
      )
  Custom implementations may be written in user code and hooked in via the
  `register_*` functions.
  """
  from collections import namedtuple
  
  import torch
  import torch.nn.functional as F
  from torch import nn
  
  from maskrcnn_benchmark.layers import FrozenBatchNorm2d
  from maskrcnn_benchmark.layers import Conv2d
  from maskrcnn_benchmark.modeling.make_layers import group_norm
  from maskrcnn_benchmark.utils.registry import Registry
  
  
  # ResNet stage specification
  StageSpec = namedtuple(
      "StageSpec",
      [
          "index",  # Index of the stage, eg 1, 2, ..,. 5
          "block_count",  # Numer of residual blocks in the stage
          "return_features",  # True => return the last feature map from this stage
      ],
  )
  
  # -----------------------------------------------------------------------------
  # Standard ResNet models
  # -----------------------------------------------------------------------------
  # ResNet-50 (including all stages)
  ResNet50StagesTo5 = tuple(
      StageSpec(index=i, block_count=c, return_features=r)
      for (i, c, r) in ((1, 3, False), (2, 4, False), (3, 6, False), (4, 3, True))
  )
  # ResNet-50 up to stage 4 (excludes stage 5)
  ResNet50StagesTo4 = tuple(
      StageSpec(index=i, block_count=c, return_features=r)
      for (i, c, r) in ((1, 3, False), (2, 4, False), (3, 6, True))
  )
  # ResNet-101 (including all stages)
  ResNet101StagesTo5 = tuple(
      StageSpec(index=i, block_count=c, return_features=r)
      for (i, c, r) in ((1, 3, False), (2, 4, False), (3, 23, False), (4, 3, True))
  )
  # ResNet-101 up to stage 4 (excludes stage 5)
  ResNet101StagesTo4 = tuple(
      StageSpec(index=i, block_count=c, return_features=r)
      for (i, c, r) in ((1, 3, False), (2, 4, False), (3, 23, True))
  )
  # ResNet-50-FPN (including all stages)
  ResNet50FPNStagesTo5 = tuple(
      StageSpec(index=i, block_count=c, return_features=r)
      for (i, c, r) in ((1, 3, True), (2, 4, True), (3, 6, True), (4, 3, True))
  )
  # ResNet-101-FPN (including all stages)
  ResNet101FPNStagesTo5 = tuple(
      StageSpec(index=i, block_count=c, return_features=r)
      for (i, c, r) in ((1, 3, True), (2, 4, True), (3, 23, True), (4, 3, True))
  )
  # ResNet-152-FPN (including all stages)
  ResNet152FPNStagesTo5 = tuple(
      StageSpec(index=i, block_count=c, return_features=r)
      for (i, c, r) in ((1, 3, True), (2, 8, True), (3, 36, True), (4, 3, True))
  )
  
  class ResNet(nn.Module):
      def __init__(self, cfg):
          super(ResNet, self).__init__()
  
          # If we want to use the cfg in forward(), then we should make a copy
          # of it and store it for later use:
          # self.cfg = cfg.clone()
  
          # Translate string names to implementations
          stem_module = _STEM_MODULES[cfg.MODEL.RESNETS.STEM_FUNC]
          stage_specs = _STAGE_SPECS[cfg.MODEL.BACKBONE.CONV_BODY]
          transformation_module = _TRANSFORMATION_MODULES[cfg.MODEL.RESNETS.TRANS_FUNC]
  
          # Construct the stem module
          self.stem = stem_module(cfg)
  
          # Constuct the specified ResNet stages
          num_groups = cfg.MODEL.RESNETS.NUM_GROUPS
          width_per_group = cfg.MODEL.RESNETS.WIDTH_PER_GROUP
          in_channels = cfg.MODEL.RESNETS.STEM_OUT_CHANNELS
          stage2_bottleneck_channels = num_groups * width_per_group
          stage2_out_channels = cfg.MODEL.RESNETS.RES2_OUT_CHANNELS
          self.stages = []
          self.return_features = {}
          for stage_spec in stage_specs:
              name = "layer" + str(stage_spec.index)
              stage2_relative_factor = 2 ** (stage_spec.index - 1)
              bottleneck_channels = stage2_bottleneck_channels * stage2_relative_factor
              out_channels = stage2_out_channels * stage2_relative_factor
              module = _make_stage(
                  transformation_module,
                  in_channels,
                  bottleneck_channels,
                  out_channels,
                  stage_spec.block_count,
                  num_groups,
                  cfg.MODEL.RESNETS.STRIDE_IN_1X1,
                  first_stride=int(stage_spec.index > 1) + 1,
              )
              in_channels = out_channels
              self.add_module(name, module)
              self.stages.append(name)
              self.return_features[name] = stage_spec.return_features
  
          # Optionally freeze (requires_grad=False) parts of the backbone
          self._freeze_backbone(cfg.MODEL.BACKBONE.FREEZE_CONV_BODY_AT)
  
      def _freeze_backbone(self, freeze_at):
          if freeze_at < 0:
              return
          for stage_index in range(freeze_at):
              if stage_index == 0:
                  m = self.stem  # stage 0 is the stem
              else:
                  m = getattr(self, "layer" + str(stage_index))
              for p in m.parameters():
                  p.requires_grad = False
  
      def forward(self, x):
          outputs = []
          x = self.stem(x)
          for stage_name in self.stages:
              x = getattr(self, stage_name)(x)
              if self.return_features[stage_name]:
                  outputs.append(x)
          return outputs
  
  
  class ResNetHead(nn.Module):
      def __init__(
          self,
          block_module,
          stages,
          num_groups=1,
          width_per_group=64,
          stride_in_1x1=True,
          stride_init=None,
          res2_out_channels=256,
          dilation=1
      ):
          super(ResNetHead, self).__init__()
  
          stage2_relative_factor = 2 ** (stages[0].index - 1)
          stage2_bottleneck_channels = num_groups * width_per_group
          out_channels = res2_out_channels * stage2_relative_factor
          in_channels = out_channels // 2
          bottleneck_channels = stage2_bottleneck_channels * stage2_relative_factor
  
          block_module = _TRANSFORMATION_MODULES[block_module]
  
          self.stages = []
          stride = stride_init
          for stage in stages:
              name = "layer" + str(stage.index)
              if not stride:
                  stride = int(stage.index > 1) + 1
              module = _make_stage(
                  block_module,
                  in_channels,
                  bottleneck_channels,
                  out_channels,
                  stage.block_count,
                  num_groups,
                  stride_in_1x1,
                  first_stride=stride,
                  dilation=dilation
              )
              stride = None
              self.add_module(name, module)
              self.stages.append(name)
  
      def forward(self, x):
          for stage in self.stages:
              x = getattr(self, stage)(x)
          return x
  
  
  def _make_stage(
      transformation_module,
      in_channels,
      bottleneck_channels,
      out_channels,
      block_count,
      num_groups,
      stride_in_1x1,
      first_stride,
      dilation=1
  ):
      blocks = []
      stride = first_stride
      for _ in range(block_count):
          blocks.append(
              transformation_module(
                  in_channels,
                  bottleneck_channels,
                  out_channels,
                  num_groups,
                  stride_in_1x1,
                  stride,
                  dilation=dilation
              )
          )
          stride = 1
          in_channels = out_channels
      return nn.Sequential(*blocks)
  
  
  class Bottleneck(nn.Module):
      def __init__(
          self,
          in_channels,
          bottleneck_channels,
          out_channels,
          num_groups,
          stride_in_1x1,
          stride,
          dilation,
          norm_func
      ):
          super(Bottleneck, self).__init__()
  
          self.downsample = None
          if in_channels != out_channels:
              down_stride = stride if dilation == 1 else 1
              self.downsample = nn.Sequential(
                  Conv2d(
                      in_channels, out_channels,
                      kernel_size=1, stride=down_stride, bias=False
                  ),
                  norm_func(out_channels),
              )
              for modules in [self.downsample,]:
                  for l in modules.modules():
                      if isinstance(l, Conv2d):
                          nn.init.kaiming_uniform_(l.weight, a=1)
  
          if dilation > 1:
              stride = 1 # reset to be 1
  
          # The original MSRA ResNet models have stride in the first 1x1 conv
          # The subsequent fb.torch.resnet and Caffe2 ResNe[X]t implementations have
          # stride in the 3x3 conv
          stride_1x1, stride_3x3 = (stride, 1) if stride_in_1x1 else (1, stride)
  
          self.conv1 = Conv2d(
              in_channels,
              bottleneck_channels,
              kernel_size=1,
              stride=stride_1x1,
              bias=False,
          )
          self.bn1 = norm_func(bottleneck_channels)
          # TODO: specify init for the above
  
          self.conv2 = Conv2d(
              bottleneck_channels,
              bottleneck_channels,
              kernel_size=3,
              stride=stride_3x3,
              padding=dilation,
              bias=False,
              groups=num_groups,
              dilation=dilation
          )
          self.bn2 = norm_func(bottleneck_channels)
  
          self.conv3 = Conv2d(
              bottleneck_channels, out_channels, kernel_size=1, bias=False
          )
          self.bn3 = norm_func(out_channels)
  
          for l in [self.conv1, self.conv2, self.conv3,]:
              nn.init.kaiming_uniform_(l.weight, a=1)
  
      def forward(self, x):
          identity = x
  
          out = self.conv1(x)
          out = self.bn1(out)
          out = F.relu_(out)
  
          out = self.conv2(out)
          out = self.bn2(out)
          out = F.relu_(out)
  
          out0 = self.conv3(out)
          out = self.bn3(out0)
  
          if self.downsample is not None:
              identity = self.downsample(x)
  
          out += identity
          out = F.relu_(out)
  
          return out
  
  
  class BaseStem(nn.Module):
      def __init__(self, cfg, norm_func):
          super(BaseStem, self).__init__()
  
          out_channels = cfg.MODEL.RESNETS.STEM_OUT_CHANNELS
  
          self.conv1 = Conv2d(
              3, out_channels, kernel_size=7, stride=2, padding=3, bias=False
          )
          self.bn1 = norm_func(out_channels)
  
          for l in [self.conv1,]:
              nn.init.kaiming_uniform_(l.weight, a=1)
  
      def forward(self, x):
          x = self.conv1(x)
          x = self.bn1(x)
          x = F.relu_(x)
          x = F.max_pool2d(x, kernel_size=3, stride=2, padding=1)
          return x
  
  
  class BottleneckWithFixedBatchNorm(Bottleneck):
      def __init__(
          self,
          in_channels,
          bottleneck_channels,
          out_channels,
          num_groups=1,
          stride_in_1x1=True,
          stride=1,
          dilation=1
      ):
          super(BottleneckWithFixedBatchNorm, self).__init__(
              in_channels=in_channels,
              bottleneck_channels=bottleneck_channels,
              out_channels=out_channels,
              num_groups=num_groups,
              stride_in_1x1=stride_in_1x1,
              stride=stride,
              dilation=dilation,
              norm_func=FrozenBatchNorm2d
          )
  
  
  class StemWithFixedBatchNorm(BaseStem):
      def __init__(self, cfg):
          super(StemWithFixedBatchNorm, self).__init__(
              cfg, norm_func=FrozenBatchNorm2d
          )
  
  
  class BottleneckWithGN(Bottleneck):
      def __init__(
          self,
          in_channels,
          bottleneck_channels,
          out_channels,
          num_groups=1,
          stride_in_1x1=True,
          stride=1,
          dilation=1
      ):
          super(BottleneckWithGN, self).__init__(
              in_channels=in_channels,
              bottleneck_channels=bottleneck_channels,
              out_channels=out_channels,
              num_groups=num_groups,
              stride_in_1x1=stride_in_1x1,
              stride=stride,
              dilation=dilation,
              norm_func=group_norm
          )
  
  
  class StemWithGN(BaseStem):
      def __init__(self, cfg):
          super(StemWithGN, self).__init__(cfg, norm_func=group_norm)
  
  
  _TRANSFORMATION_MODULES = Registry({
      "BottleneckWithFixedBatchNorm": BottleneckWithFixedBatchNorm,
      "BottleneckWithGN": BottleneckWithGN,
  })
  
  _STEM_MODULES = Registry({
      "StemWithFixedBatchNorm": StemWithFixedBatchNorm,
      "StemWithGN": StemWithGN,
  })
  
  _STAGE_SPECS = Registry({
      "R-50-C4": ResNet50StagesTo4,
      "R-50-C4-SNR": ResNet50StagesTo4,
      "R-50-C5": ResNet50StagesTo5,
      "R-101-C4": ResNet101StagesTo4,
      "R-101-C5": ResNet101StagesTo5,
      "R-50-FPN": ResNet50FPNStagesTo5,
      "R-50-FPN-RETINANET": ResNet50FPNStagesTo5,
      "R-101-FPN": ResNet101FPNStagesTo5,
      "R-101-FPN-RETINANET": ResNet101FPNStagesTo5,
      "R-152-FPN": ResNet152FPNStagesTo5,
  })
  
  
  
  
  
  
  
  
  class ChannelGate_sub(nn.Module):
      """A mini-network that generates channel-wise gates conditioned on input tensor."""
  
      def __init__(self, in_channels, num_gates=None, return_gates=False,
                   gate_activation='sigmoid', reduction=16, layer_norm=False):
          super(ChannelGate_sub, self).__init__()
          if num_gates is None:
              num_gates = in_channels
          self.return_gates = return_gates
          self.global_avgpool = nn.AdaptiveAvgPool2d(1)
          self.fc1 = nn.Conv2d(in_channels, in_channels//reduction, kernel_size=1, bias=True, padding=0)
          self.norm1 = None
          if layer_norm:
              self.norm1 = nn.LayerNorm((in_channels//reduction, 1, 1))
          self.relu = nn.ReLU(inplace=True)
          self.fc2 = nn.Conv2d(in_channels//reduction, num_gates, kernel_size=1, bias=True, padding=0)
          if gate_activation == 'sigmoid':
              self.gate_activation = nn.Sigmoid()
          elif gate_activation == 'relu':
              self.gate_activation = nn.ReLU(inplace=True)
          elif gate_activation == 'linear':
              self.gate_activation = None
          else:
              raise RuntimeError("Unknown gate activation: {}".format(gate_activation))
  
      def forward(self, x):
          input = x
          x = self.global_avgpool(x)
          x = self.fc1(x)
          if self.norm1 is not None:
              x = self.norm1(x)
          x = self.relu(x)
          x = self.fc2(x)
          if self.gate_activation is not None:
              x = self.gate_activation(x)
          if self.return_gates:
              return x
          return input * x #, input * (1 - x), x
  
  
  class ResNet_SNR(nn.Module):
  
      def __init__(self, block, layers):
          self.inplanes = 64
          super(ResNet_SNR, 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=2)
          self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
          # self.avgpool = nn.AvgPool2d(7, stride=1)
          # self.fc = nn.Linear(512 * block.expansion, num_classes)
  
          # IN bridge:
          self.IN1 = nn.InstanceNorm2d(256, affine=True)
          self.IN2 = nn.InstanceNorm2d(512, affine=True)
          self.IN3 = nn.InstanceNorm2d(1024, affine=True)
          #self.IN4 = nn.InstanceNorm2d(512, affine=True)
  
          # SE for selection:
          self.style_reid_laye1 = ChannelGate_sub(256, num_gates=256, return_gates=False,
                   gate_activation='sigmoid', reduction=16, layer_norm=False)
          self.style_reid_laye2 = ChannelGate_sub(512, num_gates=512, return_gates=False,
                   gate_activation='sigmoid', reduction=16, layer_norm=False)
          self.style_reid_laye3 = ChannelGate_sub(1024, num_gates=1024, return_gates=False,
                   gate_activation='sigmoid', reduction=16, layer_norm=False)
          # self.style_reid_laye4 = ChannelGate_sub(512, num_gates=512, return_gates=False,
          #          gate_activation='sigmoid', reduction=16, layer_norm=False)
  
          for m in self.modules():
              if isinstance(m, nn.Conv2d):
                  nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                  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.BatchNorm1d):
                  nn.init.constant_(m.weight, 1)
                  nn.init.constant_(m.bias, 0)
              elif isinstance(m, nn.InstanceNorm2d):
                  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)
                  if m.bias is not None:
                      nn.init.constant_(m.bias, 0)
  
      def bn_eval(self):
          for m in self.modules():
              if isinstance(m, nn.BatchNorm2d):
                  m.eval()
  
      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):
  
          x = self.conv1(x)
          x = self.bn1(x)
          x = self.relu(x)
          x = self.maxpool(x)
  
          x_1 = self.layer1(x)  # torch.Size([64, 256, 64, 32])
          x_1_ori = x_1
          x_IN_1 = self.IN1(x_1)
          x_style_1 = x_1 - x_IN_1
          x_style_1_reid_useful, x_style_1_reid_useless, selective_weight_useful_1 = self.style_reid_laye1(x_style_1)
          x_1 = x_IN_1 + x_style_1_reid_useful
          x_1_useless = x_IN_1 + x_style_1_reid_useless
  
  
          x_2 = self.layer2(x_1)  # torch.Size([64, 512, 32, 16])
          x_2_ori = x_2
          x_IN_2 = self.IN2(x_2)
          x_style_2 = x_2 - x_IN_2
          x_style_2_reid_useful, x_style_2_reid_useless, selective_weight_useful_2 = self.style_reid_laye2(x_style_2)
          x_2 = x_IN_2 + x_style_2_reid_useful
          x_2_useless = x_IN_2 + x_style_2_reid_useless
  
          x_3 = self.layer3(x_2)  # torch.Size([64, 1024, 16, 8])
          x_3_ori = x_3
          x_IN_3 = self.IN3(x_3)
          x_style_3 = x_3 - x_IN_3
          x_style_3_reid_useful, x_style_3_reid_useless, selective_weight_useful_3 = self.style_reid_laye3(x_style_3)
          x_3 = x_IN_3 + x_style_3_reid_useful
          x_3_useless = x_IN_3 + x_style_3_reid_useless
  
          x_4 = self.layer4(x_3)  # torch.Size([64, 2048, 16, 8])
  
          # x_4 = self.avgpool(x_4)
          # x_4 = x_4.view(x_4.size(0), -1)
          #
          # end_points['Feature'] = x_4
  
          return x_4
  
  class ResNet_SNR_v2(nn.Module):
      def __init__(self, cfg):
          super(ResNet_SNR_v2, self).__init__()
  
          # If we want to use the cfg in forward(), then we should make a copy
          # of it and store it for later use:
          # self.cfg = cfg.clone()
  
          # Translate string names to implementations
          stem_module = _STEM_MODULES[cfg.MODEL.RESNETS.STEM_FUNC]
          stage_specs = _STAGE_SPECS[cfg.MODEL.BACKBONE.CONV_BODY]
          transformation_module = _TRANSFORMATION_MODULES[cfg.MODEL.RESNETS.TRANS_FUNC]
  
          # Construct the stem module
          self.stem = stem_module(cfg)
  
          # Constuct the specified ResNet stages
          num_groups = cfg.MODEL.RESNETS.NUM_GROUPS
          width_per_group = cfg.MODEL.RESNETS.WIDTH_PER_GROUP
          in_channels = cfg.MODEL.RESNETS.STEM_OUT_CHANNELS
          stage2_bottleneck_channels = num_groups * width_per_group
          stage2_out_channels = cfg.MODEL.RESNETS.RES2_OUT_CHANNELS
          self.stages = []
          self.return_features = {}
          count = 0
          for stage_spec in stage_specs:
              count += 1
              name = "layer" + str(stage_spec.index)
              stage2_relative_factor = 2 ** (stage_spec.index - 1)
              bottleneck_channels = stage2_bottleneck_channels * stage2_relative_factor
              out_channels = stage2_out_channels * stage2_relative_factor
              module = _make_stage(
                  transformation_module,
                  in_channels,
                  bottleneck_channels,
                  out_channels,
                  stage_spec.block_count,
                  num_groups,
                  cfg.MODEL.RESNETS.STRIDE_IN_1X1,
                  first_stride=int(stage_spec.index > 1) + 1,
              )
              # SNR module:
              SNR_module = nn.Sequential(nn.InstanceNorm2d(in_channels, affine=True),
                                         ChannelGate_sub(in_channels, num_gates=in_channels, return_gates=False,
                                                         gate_activation='sigmoid', reduction=8, layer_norm=False)
                                         )
  
              in_channels = out_channels
              self.add_module(name, module)
              #if count <= 3:
              self.add_module(name+'_SNR', SNR_module)
              self.stages.append(name)
              self.return_features[name] = stage_spec.return_features
  
          # Optionally freeze (requires_grad=False) parts of the backbone
          self._freeze_backbone(cfg.MODEL.BACKBONE.FREEZE_CONV_BODY_AT)
  
          # # IN bridge:
          # self.IN1 = nn.InstanceNorm2d(256, affine=True)
          # self.IN2 = nn.InstanceNorm2d(512, affine=True)
          # self.IN3 = nn.InstanceNorm2d(1024, affine=True)
          # # self.IN4 = nn.InstanceNorm2d(512, affine=True)
          #
          # # SE for selection:
          # self.style_reid_laye1 = ChannelGate_sub(256, num_gates=256, return_gates=False,
          #          gate_activation='sigmoid', reduction=16, layer_norm=False)
          # self.style_reid_laye2 = ChannelGate_sub(512, num_gates=512, return_gates=False,
          #          gate_activation='sigmoid', reduction=16, layer_norm=False)
          # self.style_reid_laye3 = ChannelGate_sub(1024, num_gates=1024, return_gates=False,
          #          gate_activation='sigmoid', reduction=16, layer_norm=False)
          # # self.style_reid_laye4 = ChannelGate_sub(512, num_gates=512, return_gates=False,
          # #          gate_activation='sigmoid', reduction=16, layer_norm=False)
  
      def _freeze_backbone(self, freeze_at):
          if freeze_at < 0:
              return
          for stage_index in range(freeze_at):
              if stage_index == 0:
                  m = self.stem  # stage 0 is the stem
              else:
                  m = getattr(self, "layer" + str(stage_index))
              for p in m.parameters():
                  p.requires_grad = False
  
      def forward(self, x):
          outputs = []
          x = self.stem(x)
          count = 0
          for stage_name in self.stages:
              count += 1
              x = getattr(self, stage_name)(x)
              # if count == 1:
              #     x_IN_1 = self.IN1(x)
              #     x_style_1 = x - x_IN_1
              #     x_style_1_reid_useful, _, _ = self.style_reid_laye1(x_style_1)
              #     x = x_IN_1 + x_style_1_reid_useful
              # elif count == 2:
              #     x_IN_2 = self.IN2(x)
              #     x_style_2 = x - x_IN_2
              #     x_style_2_reid_useful, _, _ = self.style_reid_laye2(x_style_2)
              #     x = x_IN_2 + x_style_2_reid_useful
              # elif count == 3:
              #     x_IN_3 = self.IN3(x)
              #     x_style_3 = x - x_IN_3
              #     x_style_3_reid_useful, _, _ = self.style_reid_laye3(x_style_3)
              #     x = x_IN_3 + x_style_3_reid_useful
  
              if self.return_features[stage_name]:
                  outputs.append(x)
          return outputs