text: # Once for All: Train One Network and Specialize it for Efficient Deployment
  # Han Cai, Chuang Gan, Tianzhe Wang, Zhekai Zhang, Song Han
  # International Conference on Learning Representations (ICLR), 2020.
  
  import copy
  import random
  
  import torch
  
  from ofa.elastic_nn.modules.dynamic_layers import DynamicMBConvLayer, DynamicConvLayer, DynamicLinearLayer
  from ofa.layers import ConvLayer, IdentityLayer, LinearLayer, MBInvertedConvLayer
  from ofa.imagenet_codebase.networks.mobilenet_s4 import MobileNetS4, MobileInvertedResidualBlock
  from ofa.imagenet_codebase.utils import make_divisible, int2list
  
  
  class OFAMobileNetS4(MobileNetS4):
      # #################### Downsampled Image 저장할라고 개쓰레기같은 Static변수 만들었음. 그냥 Image 저장할 때 Image 이름 다르게할려고.
      # save_name_counter = 0  # validation forward시에 맨처음에 dummy image가 하나 들어가서 running check 하는 것 같은 작동을 보여주는데, 그래서 0으로 설정해줘서 처음 dummy image 저장한다.
  
      def __init__(self, bn_param=(0.1, 1e-5), dropout_rate=0.1, base_stage_width=None,
                   width_mult_list=1.0, ks_list=7, expand_ratio_list=6, depth_list=4, pixelshuffle_depth_list=2):
  
          self.width_mult_list = int2list(width_mult_list, 1)  # 이게 output width 조절하는 변수
          self.ks_list = int2list(ks_list, 1)
          self.expand_ratio_list = int2list(expand_ratio_list, 1)
          self.depth_list = int2list(depth_list, 1)
          self.pixelshuffle_depth_list = int2list(pixelshuffle_depth_list, 1)
          self.base_stage_width = base_stage_width
  
          self.width_mult_list.sort()
          self.ks_list.sort()
          self.expand_ratio_list.sort()
          self.depth_list.sort()
          self.pixelshuffle_depth_list.sort()
                      # FROM [3,      64, 64, 64, 64,     64, 64,     64,  64]
          base_stage_width = [64,     64, 64, 64, 64,     64, 64,     256, 3]
                           # [        ResBlock              ResCon    Shu]
                           # [1,      4,  4,  4,  4,      1,  1,      2,   1]
                           # [  Skip,                       Con]
  
          # final_expand_width = [
          #     make_divisible(base_stage_width[-2] * max(self.width_mult_list), 8) for _ in self.width_mult_list
          # ]
          # last_channel = [
          #     make_divisible(base_stage_width[-1] * max(self.width_mult_list), 8) for _ in self.width_mult_list
          # ]
  
          stride_stages = [1,     1, 1, 1, 1,     1, 1,     1, 1]
          act_stages = [None,   'relu6', 'relu6', 'relu6', 'relu6',     None, None,     'pixelshuffle', None]
          se_stages = [False,     False, False, False, False,     False, False,     False, False]
          if depth_list is None:
              n_block_list = [1, 2, 3, 4, 2, 3]
              self.depth_list = [4, 4]
              print('Use MobileNetV3 Depth Setting')
          else:
              n_block_list = [1] + [max(self.depth_list)]*4 + [1]*2 + [max(self.pixelshuffle_depth_list)] + [1]  # 2는 pixelshuffle, pixelunshuffle의 depth
              # [1, 4, 4, 4, 4, 1, 1, 2, 1]
          width_list = []
          for base_width in base_stage_width:
              # width = [make_divisible(base_width * width_mult, 8) for width_mult in self.width_mult_list]
              width = [make_divisible(base_width * width_mult, 1) for width_mult in self.width_mult_list]
              width_list.append(width)
  
          #################################################################################################### decoder first conv block
          dec_first_conv_block = ConvLayer(3, max(width_list[0]), kernel_size=5, stride=stride_stages[0], act_func=act_stages[0], use_bn=True)
  
          #################################################################################################### decoder inverted residual blocks
          self.block_group_info = []
          blocks = []
          _block_index = 0
          feature_dim = width_list[0]
  
          for width, n_block, s, act_func, use_se in zip(width_list[1:5], n_block_list[1:5],
                                                         stride_stages[1:5], act_stages[1:5], se_stages[1:5]):
              self.block_group_info.append([_block_index + i for i in range(n_block)])
              _block_index += n_block
  
              output_channel = width
              for i in range(n_block):
                  if i == 0:
                      stride = s
                  else:
                      stride = 1
                  mobile_inverted_conv = DynamicMBConvLayer(
                      in_channel_list=feature_dim, out_channel_list=output_channel, kernel_size_list=ks_list,
                      expand_ratio_list=expand_ratio_list, stride=stride, act_func=act_func, use_se=use_se,
                  )
                  shortcut = IdentityLayer(feature_dim, feature_dim)
                  blocks.append(MobileInvertedResidualBlock(mobile_inverted_conv, shortcut))
                  feature_dim = output_channel
  
          #################################################################################################### decoder final conv blocks
          dec_final_conv_blocks = []
          for width, n_block, s, act_func, use_se in zip(width_list[5:7], n_block_list[5:7],
                                                         stride_stages[5:7], act_stages[5:7], se_stages[5:7]):
              # self.block_group_info.append([_block_index + i for i in range(n_block)])
              # _block_index += n_block
  
              output_channel = width
              for i in range(n_block):
                  if i == 0:
                      stride = s
                  else:
                      stride = 1
                  dec_final_conv_blocks.append(ConvLayer(max(feature_dim), max(output_channel), kernel_size=5, stride=s, act_func=act_func, use_bn=True))
                  feature_dim = output_channel
  
          #################################################################################################### decoder shuffle
      # for width, n_block, s, act_func, use_se in zip(width_list[11], n_block_list[11],
      #                                                 stride_stages[11], act_stages[11], se_stages[11]):
          self.block_group_info.append([_block_index + i for i in range(n_block_list[7])])
          _block_index += n_block_list[7]
  
          output_channel = width_list[7]
          for i in range(n_block_list[7]):
              if i == 0:
                  stride = stride_stages[7]
              else:
                  stride = 1
              blocks.append(ConvLayer(max(feature_dim), max(output_channel), kernel_size=5, stride=s, act_func=act_stages[7], use_bn=True))
  
          #################################################################################################### decoder final output conv block
          dec_final_output_conv_block = ConvLayer(max(feature_dim), max(width_list[8]), kernel_size=5, stride=stride_stages[8], act_func=act_stages[8], use_bn=True)
  
          ####################################################################################################
          # runtime_depth
          self.runtime_depth = [len(block_idx) for block_idx in self.block_group_info]
  
          super(OFAMobileNetS4, self).__init__(blocks,
                                              dec_first_conv_block, dec_final_conv_blocks, dec_final_output_conv_block, self.runtime_depth)
  
          # set bn param
          self.set_bn_param(momentum=bn_param[0], eps=bn_param[1])
  
      """ MyNetwork required methods """
  
      @staticmethod
      def name():
          return 'OFAMobileNetS4'
  
  
      def forward(self, x):
          ######### decoder first conv block
          x = self.dec_first_conv_block(x)
  
          dec_big_skip = x
  
          ########## decoder inverted residual blocks
          for stage_id, block_idx in enumerate(self.block_group_info[:4]):
              depth = self.runtime_depth[stage_id]
              active_idx = block_idx[:depth]
              for idx in active_idx:
                  x = self.blocks[idx](x)
  
          ########## decoder final conv blocks
          for i, dec_final_conv_block in enumerate(self.dec_final_conv_blocks):
              x = dec_final_conv_block(x)
              if i == 0:
                  x += dec_big_skip
  
          ########## decoder shuffle
          for stage_id, block_idx in enumerate(self.block_group_info[4:]):
              depth = self.runtime_depth[stage_id]
              active_idx = block_idx[:depth]
              for idx in active_idx:
                  x = self.blocks[idx](x)
  
          ########## decoder final output conv block
          x = self.dec_final_output_conv_block(x)
  
          #################### For saving downsampled image. 필요할 때만 주석 해제해서 사용하면 된다.
          # im = tensor2img_np(x)
          # # print(OFAMobileNetX4.save_name_counter)
          # im = Image.fromarray(im)
          # im.save("/SSD/uvg-1080p/readysetgo_120fps_448center_4xLearned/%04d.png" % OFAMobileNetX4.save_name_counter)
          # OFAMobileNetX4.save_name_counter += 1
  
          return x
          
  
      @property
      def module_str(self):
          _str = ''
          for stage_id, block_idx in enumerate(self.block_group_info):
              depth = self.runtime_depth[stage_id]
              active_idx = block_idx[:depth]
              for idx in active_idx:
                  _str += self.blocks[idx].module_str + '\n'
  
          _str += self.dec_first_conv_block.module_str + '\n'
  
          for block in self.dec_final_conv_blocks:
              _str += block.module_str + '\n'
          
          _str += self.dec_final_output_conv_block.module_str + '\n'
  
          return _str
  
  
      @property
      def config(self):
          return {
              'name': OFAMobileNetS4.__name__,
              'bn': self.get_bn_params(),
              'blocks': [
                  block.config for block in self.blocks
              ],
              'dec_first_conv_block': self.dec_first_conv_block.config,
              'dec_final_conv_blocks': [
                  block.config for block in self.dec_final_conv_blocks
              ],
              'dec_final_output_conv_block': self.dec_final_output_conv_block.config,
          }
  
  
  #################################################################################################### 아래부분 일단 Teacher Training에서는 사용되지 않음
      @staticmethod
      def build_from_config(config):
          raise ValueError('do not support this function')
  
      def load_weights_from_net(self, src_model_dict):
          model_dict = self.state_dict()
          # for key, value in model_dict.items() :
          #     print(key)
          ########## 왜인지는 모르겠지만 module. 이 붙어있는 state dict 때문에...
          module_flag = False
          for key, value in src_model_dict.items():
              if key.find("module.") != -1:
                  module_flag = True
  
          for key in src_model_dict:
              ########## 왜인지는 모르겠지만 module. 이 붙어있는 state dict 때문에...
              if key.find("module.") != -1:
                  key = key.replace("module.", '')
  
              if key in model_dict:
                  new_key = key
              elif '.bn.bn.' in key:
                  new_key = key.replace('.bn.bn.', '.bn.')
              elif '.conv.conv.weight' in key:
                  new_key = key.replace('.conv.conv.weight', '.conv.weight')
              elif '.linear.linear.' in key:
                  new_key = key.replace('.linear.linear.', '.linear.')
              ##############################################################################
              elif '.linear.' in key:
                  new_key = key.replace('.linear.', '.linear.linear.')
              elif 'bn.' in key:
                  new_key = key.replace('bn.', 'bn.bn.')
              elif 'conv.weight' in key:
                  new_key = key.replace('conv.weight', 'conv.conv.weight')
              else:
                  raise ValueError(key)
              assert new_key in model_dict, '%s' % new_key
              ########## 왜인지는 모르겠지만 module. 이 붙어있는 state dict 때문에...
              if module_flag:
                  model_dict[new_key] = src_model_dict["module."+key]
              else:
                  model_dict[new_key] = src_model_dict[key]
          self.load_state_dict(model_dict)
  
      """ set, sample and get active sub-networks """
  
      def set_active_subnet(self, wid=None, ks=None, e=None, d=None, pixel_d=None):
          # width_mult_id = int2list(wid, 4 + len(self.block_group_info))
          # ks = int2list(ks, len(self.blocks) - 1)
          # expand_ratio = int2list(e, len(self.blocks) - 1)
          # depth = int2list(d, len(self.block_group_info) - 2)
  
          # for block, k, e in zip(self.blocks[1:], ks, expand_ratio):
          #     if k is not None:
          #         block.mobile_inverted_conv.active_kernel_size = k
          #     if e is not None:
          #         block.mobile_inverted_conv.active_expand_ratio = e
  
          # for i, d in enumerate(depth):
          #     if d is not None:
          #         self.runtime_depth[i] = min(len(self.block_group_info[i]), d)
          ks = int2list(ks, len(self.blocks) - 2)
          expand_ratio = int2list(e, len(self.blocks) - 2)
          depth = int2list(d, len(self.block_group_info) - 1)
          pixelshuffle_depth = int2list(pixel_d, 1)
  
          depth.insert(-1, pixelshuffle_depth[0])
  
          for block, k, e in zip(self.blocks[:-1], ks, expand_ratio):
              if k is not None:
                  block.mobile_inverted_conv.active_kernel_size = k
              if e is not None:
                  block.mobile_inverted_conv.active_expand_ratio = e
          
          for i, d in enumerate(depth):
              if d is not None:
                  self.runtime_depth[i] = min(len(self.block_group_info[i]), d)
  
      def set_constraint(self, include_list, constraint_type='depth'):
          if constraint_type == 'depth':
              self.__dict__['_depth_include_list'] = include_list.copy()
          elif constraint_type == 'expand_ratio':
              self.__dict__['_expand_include_list'] = include_list.copy()
          elif constraint_type == 'kernel_size':
              self.__dict__['_ks_include_list'] = include_list.copy()
          elif constraint_type == 'width_mult':
              self.__dict__['_width_mult_include_list'] = include_list.copy()
          elif constraint_type == 'pixelshuffle_depth':
              self.__dict__['_pixelshuffle_depth_include_list'] = include_list.copy()
          else:
              raise NotImplementedError
  
      def clear_constraint(self):
          self.__dict__['_depth_include_list'] = None
          self.__dict__['_expand_include_list'] = None
          self.__dict__['_ks_include_list'] = None
          self.__dict__['_width_mult_include_list'] = None
          self.__dict__['_pixelshuffle_depth_include_list'] = None
  
      def sample_active_subnet(self):
          # self.blocks, self.block_group_info 두 개만 shrinking에 포함되는 block들에 대한 정보를 가지고 있다 
          ks_candidates = self.ks_list if self.__dict__.get('_ks_include_list', None) is None \
              else self.__dict__['_ks_include_list']
          expand_candidates = self.expand_ratio_list if self.__dict__.get('_expand_include_list', None) is None \
              else self.__dict__['_expand_include_list']
          depth_candidates = self.depth_list if self.__dict__.get('_depth_include_list', None) is None else \
              self.__dict__['_depth_include_list']
          pixelshuffle_depth_candidates = self.pixelshuffle_depth_list if self.__dict__.get('_pixelshuffle_depth_include_list', None) is None else \
              self.__dict__['_pixelshuffle_depth_include_list']
  
          # sample width_mult
          width_mult_setting = None
  
          # sample kernel size
          ks_setting = []
          if not isinstance(ks_candidates[0], list):
              ks_candidates = [ks_candidates for _ in range(len(self.blocks) - 2)]  # (length of self.blocks) - (number of pixelshuffles)
          for k_set in ks_candidates:
              k = random.choice(k_set)
              ks_setting.append(k)
  
          # sample expand ratio
          expand_setting = []
          if not isinstance(expand_candidates[0], list):
              expand_candidates = [expand_candidates for _ in range(len(self.blocks) - 2)]  # (length of self.blocks) - (number of pixelshuffles)
          for e_set in expand_candidates:
              e = random.choice(e_set)
              expand_setting.append(e)
  
          # sample depth
          depth_setting = []
          if not isinstance(depth_candidates[0], list):
              depth_candidates = [depth_candidates for _ in range(len(self.block_group_info) - 1)]  # (length of self.block_group_info) - (number of pixelshuffle groups)
          for d_set in depth_candidates:
              d = random.choice(d_set)
              depth_setting.append(d)
          
          # sample pixelshuffle depth
          pixelshuffle_depth_setting = []
          if not isinstance(pixelshuffle_depth_candidates[0], list):
              pixelshuffle_depth_candidates = [pixelshuffle_depth_candidates for _ in range(1)]  # (number of pixelshuffle groups) -> 1 because of mirroring
          for pixel_d_set in pixelshuffle_depth_candidates:
              pixel_d = random.choice(pixel_d_set)
              pixelshuffle_depth_setting.append(pixel_d)
  
          self.set_active_subnet(width_mult_setting, ks_setting, expand_setting, depth_setting, pixelshuffle_depth_setting)
  
          return {
              'wid': width_mult_setting,
              'ks': ks_setting,
              'e': expand_setting,
              'd': depth_setting,
              'pixel_d': pixelshuffle_depth_setting,
          }
  
  #################################################################################################### 이 부분은 Network Specific하게 뽑아내서 Inference 할 때, 즉 굳이 사용할 필요없다
      def get_active_subnet(self, preserve_weight=True):
          first_conv = copy.deepcopy(self.first_conv)
          blocks = [copy.deepcopy(self.blocks[0])]
  
          final_expand_layer = copy.deepcopy(self.final_expand_layer)
          feature_mix_layer = copy.deepcopy(self.feature_mix_layer)
          classifier = copy.deepcopy(self.classifier)
  
          input_channel = blocks[0].mobile_inverted_conv.out_channels
          # blocks
          for stage_id, block_idx in enumerate(self.block_group_info):
              depth = self.runtime_depth[stage_id]
              active_idx = block_idx[:depth]
              stage_blocks = []
              for idx in active_idx:
                  stage_blocks.append(MobileInvertedResidualBlock(
                      self.blocks[idx].mobile_inverted_conv.get_active_subnet(input_channel, preserve_weight),
                      copy.deepcopy(self.blocks[idx].shortcut)
                  ))
                  input_channel = stage_blocks[-1].mobile_inverted_conv.out_channels
              blocks += stage_blocks
  
          _subnet = MobileNetV3(first_conv, blocks, final_expand_layer, feature_mix_layer, classifier)
          _subnet.set_bn_param(**self.get_bn_param())
          return _subnet
  
      def get_active_net_config(self):
          # first conv
          first_conv_config = self.first_conv.config
          first_block_config = self.blocks[0].config
          if isinstance(self.first_conv, DynamicConvLayer):
              first_conv_config = self.first_conv.get_active_subnet_config(3)
              first_block_config = {
                  'name': MobileInvertedResidualBlock.__name__,
                  'mobile_inverted_conv': self.blocks[0].mobile_inverted_conv.get_active_subnet_config(
                      first_conv_config['out_channels']
                  ),
                  'shortcut': self.blocks[0].shortcut.config if self.blocks[0].shortcut is not None else None,
              }
          final_expand_config = self.final_expand_layer.config
          feature_mix_layer_config = self.feature_mix_layer.config
          if isinstance(self.final_expand_layer, DynamicConvLayer):
              final_expand_config = self.final_expand_layer.get_active_subnet_config(
                  self.blocks[-1].mobile_inverted_conv.active_out_channel)
              feature_mix_layer_config = self.feature_mix_layer.get_active_subnet_config(
                  final_expand_config['out_channels'])
          classifier_config = self.classifier.config
          if isinstance(self.classifier, DynamicLinearLayer):
              classifier_config = self.classifier.get_active_subnet_config(self.feature_mix_layer.active_out_channel)
  
          block_config_list = [first_block_config]
          input_channel = first_block_config['mobile_inverted_conv']['out_channels']
          for stage_id, block_idx in enumerate(self.block_group_info):
              depth = self.runtime_depth[stage_id]
              active_idx = block_idx[:depth]
              stage_blocks = []
              for idx in active_idx:
                  middle_channel = make_divisible(round(input_channel *
                                                        self.blocks[idx].mobile_inverted_conv.active_expand_ratio), 8)
                  stage_blocks.append({
                      'name': MobileInvertedResidualBlock.__name__,
                      'mobile_inverted_conv': {
                          'name': MBInvertedConvLayer.__name__,
                          'in_channels': input_channel,
                          'out_channels': self.blocks[idx].mobile_inverted_conv.active_out_channel,
                          'kernel_size': self.blocks[idx].mobile_inverted_conv.active_kernel_size,
                          'stride': self.blocks[idx].mobile_inverted_conv.stride,
                          'expand_ratio': self.blocks[idx].mobile_inverted_conv.active_expand_ratio,
                          'mid_channels': middle_channel,
                          'act_func': self.blocks[idx].mobile_inverted_conv.act_func,
                          'use_se': self.blocks[idx].mobile_inverted_conv.use_se,
                      },
                      'shortcut': self.blocks[idx].shortcut.config if self.blocks[idx].shortcut is not None else None,
                  })
                  input_channel = self.blocks[idx].mobile_inverted_conv.active_out_channel
              block_config_list += stage_blocks
  
          return {
              'name': MobileNetV3.__name__,
              'bn': self.get_bn_param(),
              'first_conv': first_conv_config,
              'blocks': block_config_list,
              'final_expand_layer': final_expand_config,
              'feature_mix_layer': feature_mix_layer_config,
              'classifier': classifier_config,
          }
  
      """ Width Related Methods """
  
      def re_organize_middle_weights(self, expand_ratio_stage=0):
          for block in self.blocks[2:-2]:
              block.mobile_inverted_conv.re_organize_middle_weights(expand_ratio_stage)
  
  
  ########## For visualize and save learned downsampling image
  import math
  from PIL import Image
  import numpy as np
  import torch
  from torchvision.utils import make_grid
  
  """
  Converts a Tensor into an image Numpy array
  Input should be either in 4D(B,(3/1),H,W), 3D(C,H,W), or 2D(H,W)
  If input is in 4D, it is splited along the first dimension to provide grid view.
  Otherwise, the tensor is assume to be single image.
  Input type: float [-1, 1] (default)
  Output type: np.uint8 [0,255] (default)
  Output dim: 3D(H,W,C) (for 4D and 3D input) or 2D(H,W) (for 2D input)
  """
  def tensor2img_np(tensor, out_type=np.uint8, min_max=(0,1)):
      tensor = tensor.float().cpu().clamp_(*min_max) # Clamp is for on hard_tanh
      tensor = (tensor - min_max[0]) / (min_max[1] - min_max[0])
      n_dim = tensor.dim()
      if n_dim == 4:
          n_img = len(tensor)
          img_np = make_grid(tensor, nrow=int(math.sqrt(n_img)), normalize=False).detach().numpy()
          img_np = np.transpose(img_np, (1, 2, 0))
      elif n_dim == 3:
          img_np = tensor.numpy()
          img_np = np.transpose(img_np, (1, 2, 0))
      elif n_dim == 2:
          img_np = tensor.numpy()
      else:
          raise TypeError('Only support 4D, 3D and 2D tensor. But receieved tensor with dimension = %d' % n_dim)
      if out_type == np.uint8:
          img_np = (img_np * 255.0).round() # This is important. Unlike matlab, numpy.unit8() WILL NOT round by default.
      return img_np.astype(out_type)
  
  def rgb2gray(img):
      in_img_type = img.dtype
      img.astype(np.float64)
      img_gray = np.dot(img[...,:3], [0.299, 0.587, 0.114]).round()
      return img_gray.astype(in_img_type)
  
  def rgb2y(img):
      assert(img.dtype == np.uint8)
      in_img_type = img.dtype
      img.astype(np.float64)
      img_y = ((np.dot(img[...,:3], [65.481, 128.553, 24.966])) / 255.0 + 16.0).round()
      return img_y.astype(in_img_type)