text: # -*- coding:utf-8 -*-
  
  # Copyright (C) 2020. Huawei Technologies Co., Ltd. All rights reserved.
  # This program is free software; you can redistribute it and/or modify
  # it under the terms of the MIT License.
  # This program is distributed in the hope that it will be useful,
  # but WITHOUT ANY WARRANTY; without even the implied warranty of
  # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  # MIT License for more details.
  
  """Some function need in spnet.py."""
  
  import torch.nn as nn
  import os.path as osp
  import torch
  from collections import OrderedDict
  from terminaltables import AsciiTable
  from mmcv.runner import get_dist_info
  
  
  def dirac_init(module, bias=0):
      """Load state_dict to a module. This method is modified from :meth:`torch.nn.Module.load_state_dict`.
  
      :param module: Module that receives the state_dict
      :type module: OrderedDict
      :param bias: bias
      :type bias: int
      """
      nn.init.dirac_(module.weight)
      if hasattr(module, 'bias') and module.bias is not None:
          nn.init.constant_(module.bias, bias)
  
  
  def load_state_dict(module, state_dict, mapping, logger=None, mb_mapping=None):  # noqa: C901
      """Load state_dict to a module. This method is modified from :meth:`torch.nn.Module.load_state_dict`.
  
      :param module: Module that receives the state_dict
      :type module: OrderedDict
      :param state_dict: Weights
      :type state_dict: Module
      :param mapping: serail backbone with related pretrained weight name
      :type mapping: dict
      :param logger: Logger to log the error
      :type logger: logging.Logger
      :param mb_mapping: parallel backbone with related serial backbone
      :type mb_mapping: dict
      """
      unexpected_keys = []
      shape_mismatch_pairs = []
      mb_keys = []
  
      own_state = module.state_dict()
      for name, param in state_dict.items():
          if name.find('mb') >= 0:
              continue
          if name not in own_state and name not in mapping.keys():
              unexpected_keys.append(name)
              continue
          if name in mapping.keys():
              if mapping[name] not in own_state:
                  unexpected_keys.append(name)
                  continue
              name = mapping[name]
          if isinstance(param, torch.nn.Parameter):
              # backwards compatibility for serialized parameters
              param = param.data
          if param.size() != own_state[name].size():
              shape_mismatch_pairs.append(
                  [name, own_state[name].size(),
                   param.size()])
              continue
          own_state[name].copy_(param)
          if mb_mapping is not None and name in mb_mapping.keys():
              for mb_name in mb_mapping[name]:
                  own_state[mb_name].copy_(param)
                  mb_keys.append(mb_name)
  
      all_missing_keys = set(own_state.keys()) - set(state_dict.keys()) - set(mapping.values()) - set(mb_keys)
      unexpected_keys = set(unexpected_keys) - set(mb_keys)
      # ignore "num_batches_tracked" of BN layers
      missing_keys = [
          key for key in all_missing_keys if 'num_batches_tracked' not in key
      ]
  
      rank, _ = get_dist_info()
      print_err = ['The model and loaded state dict do not match exactly\n']
      if rank == 0:
          if unexpected_keys:
              print_err.append('unexpected key in source state_dict: {}\n'.format(
                  ', '.join(unexpected_keys)))
          if missing_keys:
              print_err.append('missing keys in source state_dict: {}\n'.format(
                  ', '.join(missing_keys)))
          if shape_mismatch_pairs:
              print_err.append('mismatched shape keys: {}\n'.format(
                  ', '.join(shape_mismatch_pairs)))
          if logger is not None:
              logger.warning(print_err)
          else:
              print(print_err)
  
  
  def load_checkpoint(model,
                      filename,
                      pretrain_to_own,
                      map_location='cpu',
                      logger=None,
                      mb_mapping=None):
      """Load checkpoint from a pretrained checkpoint file.
  
      :param module: Module that receives the state_dict
      :type module: OrderedDict
      :param filename: A filepath.
      :type filename: str
      :param pretrain_to_own: serail backbone with related pretrained weight name
      :type pretrain_to_own: dict
      :param logger: Logger to log the error
      :type logger: logging.Logger
      :param mb_mapping: parallel backbone with related serial backbone
      :type mb_mapping: dict
      :param map_location: map location
      :type map_location: str
      :return: checkpoint
      :rtype: checkpoint
      """
      # load checkpoint from file
      if not osp.isfile(filename):
          raise IOError('{} is not a checkpoint file'.format(filename))
      checkpoint = torch.load(filename, map_location=map_location)
      # get state_dict from checkpoint
      if isinstance(checkpoint, OrderedDict):
          state_dict = checkpoint
      elif isinstance(checkpoint, dict) and 'state_dict' in checkpoint:
          state_dict = checkpoint['state_dict']
      else:
          raise RuntimeError(
              'No state_dict found in checkpoint file {}'.format(filename))
      # strip prefix of state_dict
      if list(state_dict.keys())[0].startswith('module.'):
          state_dict = {k[7:]: v for k, v in checkpoint['state_dict'].items()}
      # load state_dict
      if hasattr(model, 'module'):
          load_state_dict(model.module, state_dict, pretrain_to_own, logger, mb_mapping)
      else:
          load_state_dict(model, state_dict, pretrain_to_own, logger, mb_mapping)
      return checkpoint
  
  
  def remove_layers(layer_arch, p_arch, len_dis):
      """Remove expand layers in serailnet than pretrained architecture.
  
      :param layer_arch: current serailnet encode
      :type layer_arch: str
      :param p_arch: pretrained serailnet encode
      :type p_arch: str
      :param len_dis: num of blocks more than pretrain
      :type len_dis: int
      :return: more expand block name
      :rtype: list
      """
      def map_index_to_name(index, layer_arch):
          name = []
          for i in index:
              for l in range(1, len(layer_arch) + 1):
                  if i < len(''.join(layer_arch[:l])):
                      block_i = i - len(''.join(layer_arch[:l - 1]))
                      n = 'layer' + str(l) + '.' + str(block_i) + '.'
                      name.append(n)
                      break
          return name
  
      def find_diff_index(pos, new_arch, old_arch):
          agg_arch = zip(new_arch, old_arch)
          for i, tup in enumerate(agg_arch):
              if tup[0] != tup[1]:
                  if len(pos) < 1:
                      pos.append(i)
                  else:
                      pos.append(i + pos[-1] + 1)
                  return find_diff_index(pos, new_arch[i + 1:], old_arch[i:])
                  break
          return pos
  
      def get_bn_layers(remove_name, layer_arch, p_arch):
          bn_layers = [n + 'bn{}'.format(i) for n in remove_name for i in range(1, 4)]
          p_len = len(p_arch.split('-'))
          for i in range(1, len(layer_arch) - p_len + 1):
              bn_layers.append('layer{}.0.downsample.1'.format(p_len + i))
          return bn_layers
  
      pretrained_arch = ''.join(p_arch.split('-'))
      remove_index = find_diff_index([], ''.join(layer_arch), pretrained_arch)
      if len(remove_index) < len_dis:
          if len(layer_arch) > len(p_arch.split('-')):
              len_dis = len_dis - (len(layer_arch) - len(p_arch.split('-')))
          remove_index = remove_index + list(range(len(''.join(layer_arch)) - len_dis + len(remove_index),
                                                   len(''.join(layer_arch))))
      remove_name = map_index_to_name(remove_index, layer_arch)
      bn_layers = get_bn_layers(remove_name, layer_arch, p_arch)
      return remove_name, bn_layers
  
  
  def match_name(own_names, checkpoint_names):
      """Generate the mapping dict between current serialnet layer and pretrain.
  
      :param own_names: keys of current serialnet
      :type own_names: list
      :param checkpoint_names: keys of pretrained serailnet
      :type checkpoint_names: list
      :return: serail backbone with related pretrained weight name
      :rtype: dict
      """
      pretrain_to_own = dict()
      print("matching pretrained model with new distributed architecture....")
      # build mapping1: blocks id to own name
      bid_to_own = []
      for name in own_names:
          if name.find('layer') >= 0:
              block_name = '.'.join(name.split('.')[:2])
              if block_name not in bid_to_own:
                  bid_to_own.append(block_name)
  
      # map pretrained name to own name
      flag = ''
      bid = -1
      for name in checkpoint_names:
          name_copy = name
          if 'backbone' in name_copy:
              name_copy = name_copy.replace('backbone.', '')
              pretrain_to_own[name] = name_copy
          if name_copy.find('layer') >= 0:
              temp = name_copy.split('.')
              if temp[0] != flag:
                  s = bid + 1
              flag = temp[0]
              bid = s + int(temp[1])
              pretrain_to_own[name] = '.'.join([bid_to_own[bid]] + temp[2:])
  
      return pretrain_to_own
  