text: import logging
  import os
  
  import torch
  import torch.nn as nn
  from mmcv.runner import master_only
  from torch.nn.parallel import DataParallel, DistributedDataParallel
  
  import mmsr.models.lr_scheduler as lr_scheduler
  
  logger = logging.getLogger('base')
  
  
  class BaseModel():
      """Base model.
      """
  
      def __init__(self, opt):
          if opt['dist']:
              self.rank = torch.distributed.get_rank()
          else:
              self.rank = -1  # non-dist training
  
          self.opt = opt
          self.device = torch.device(
              'cuda' if opt['gpu_ids'] is not None else 'cpu')
          self.is_train = opt['is_train']
          self.schedulers = []
          self.optimizers = []
  
      def feed_data(self, data):
          pass
  
      def optimize_parameters(self):
          pass
  
      def get_current_visuals(self):
          pass
  
      def save(self, epoch, current_iter):
          """Save networks and training state."""
          pass
  
      def validation(self, dataloader, current_iter, tb_logger, save_img=False):
          """Validation function.
  
          Args:
              dataloder (torch.utils.data.DataLoader): Validation dataloader.
              current_iter (int): Current iteration.
              tb_logger (tensorboard logger): Tensorboard logger.
              save_img (bool): Whether to save images. Default: False.
          """
          if self.opt['dist']:
              self.dist_validation(dataloader, current_iter, tb_logger, save_img)
          else:
              self.nondist_validation(dataloader, current_iter, tb_logger,
                                      save_img)
  
      def get_current_log(self):
          return self.log_dict
  
      def model_to_device(self, net):
          """Model to device. It also warps models with DistributedDataParallel
          or DataParallel.
  
          Args:
              net (nn.Module)
          """
          net = net.to(self.device)
          if self.opt['dist']:
              net = DistributedDataParallel(
                  net, device_ids=[torch.cuda.current_device()])
          else:
              net = DataParallel(net)
          return net
  
      def setup_schedulers(self):
          """Set up schedulers."""
          train_opt = self.opt['train']
          if train_opt['lr_scheme'] == 'MultiStepLR':
              for optimizer in self.optimizers:
                  self.schedulers.append(
                      lr_scheduler.MultiStepRestartLR(
                          optimizer,
                          train_opt['lr_steps'],
                          gamma=train_opt['lr_gamma']))
          elif train_opt['lr_scheme'] == 'MultiStepRestartLR':
              for optimizer in self.optimizers:
                  self.schedulers.append(
                      lr_scheduler.MultiStepRestartLR(
                          optimizer,
                          train_opt['lr_steps'],
                          restarts=train_opt['restarts'],
                          restart_weights=train_opt['restart_weights'],
                          gamma=train_opt['lr_gamma']))
          elif train_opt['lr_scheme'] == 'CosineAnnealingRestartLR':
              for optimizer in self.optimizers:
                  self.schedulers.append(
                      lr_scheduler.CosineAnnealingRestartLR(
                          optimizer,
                          train_opt['period'],
                          eta_min=train_opt['eta_min'],
                          restarts=train_opt['restarts'],
                          restart_weights=train_opt['restart_weights']))
          else:
              raise NotImplementedError(
                  f"{train_opt['lr_scheme']} learning rate scheme is not "
                  'implemented.')
  
      def _get_network_description(self, net):
          """Get the string and total parameters of the network"""
          if isinstance(net, nn.DataParallel) or isinstance(
                  net, DistributedDataParallel):
              net = net.module
          return str(net), sum(map(lambda x: x.numel(), net.parameters()))
  
      @master_only
      def print_network(self, net):
          """Print the str and parameter number of a network.
  
          Args:
              net (nn.Module)
          """
          net_str, net_params = self._get_network_description(net)
          if isinstance(net, nn.DataParallel) or isinstance(
                  net, DistributedDataParallel):
              net_cls_str = (f'{net.__class__.__name__} - '
                             f'{net.module.__class__.__name__}')
          else:
              net_cls_str = f'{net.__class__.__name__}'
  
          logger.info(
              f'Network: {net_cls_str}, with parameters: {net_params:,d}')
          logger.info(net_str)
  
      def _set_lr(self, lr_groups_l):
          """Set learning rate for warmup.
  
          Args:
              lr_groups_l (list): List for lr_groups, each for an optimizer.
          """
          for optimizer, lr_groups in zip(self.optimizers, lr_groups_l):
              for param_group, lr in zip(optimizer.param_groups, lr_groups):
                  param_group['lr'] = lr
  
      def _get_init_lr(self):
          """Get the initial lr, which is set by the scheduler.
          """
          init_lr_groups_l = []
          for optimizer in self.optimizers:
              init_lr_groups_l.append(
                  [v['initial_lr'] for v in optimizer.param_groups])
          return init_lr_groups_l
  
      def update_learning_rate(self, current_iter, warmup_iter=-1):
          """Update learning rate.
  
          Args:
              current_iter (int): Current iteration.
              warmup_iter (int)： Warmup iter numbers. -1 for no warmup.
                  Default： -1.
          """
          if current_iter > 1:
              for scheduler in self.schedulers:
                  scheduler.step()
          # set up warm-up learning rate
          if current_iter < warmup_iter:
              # get initial lr for each group
              init_lr_g_l = self._get_init_lr()
              # modify warming-up learning rates
              # currently only support linearly warm up
              warm_up_lr_l = []
              for init_lr_g in init_lr_g_l:
                  warm_up_lr_l.append(
                      [v / warmup_iter * current_iter for v in init_lr_g])
              # set learning rate
              self._set_lr(warm_up_lr_l)
  
      def get_current_learning_rate(self):
          return [
              param_group['lr']
              for param_group in self.optimizers[0].param_groups
          ]
  
      @master_only
      def save_network(self, net, net_label, current_iter):
          """Save networks.
  
          Args:
              net (nn.Module): Network to be saved.
              net_label (str): Network label.
              current_iter (int): Current iter number.
          """
          if current_iter == -1:
              current_iter = 'latest'
          save_filename = f'{net_label}_{current_iter}.pth'
          save_path = os.path.join(self.opt['path']['models'], save_filename)
          if isinstance(net, nn.DataParallel) or isinstance(
                  net, DistributedDataParallel):
              net = net.module
          state_dict = net.state_dict()
          for key, param in state_dict.items():
              if key.startswith('module.'):  # remove unnecessary 'module.'
                  key = key[7:]
              state_dict[key] = param.cpu()
          torch.save(state_dict, save_path)
  
      def _print_different_keys_loading(self, crt_net, load_net, strict=True):
          """Print keys with differnet name or different size when loading models.
  
          1. Print keys with differnet names.
          2. If strict=False, print the same key but with different tensor size.
              It also ignore these keys with different sizes.
  
          Args:
              crt_net (torch model): Current network.
              load_net (dict): Loaded network.
              strict (bool): Whether strictly loaded. Default: True.
          """
          if isinstance(crt_net, nn.DataParallel) or isinstance(
                  crt_net, DistributedDataParallel):
              crt_net = crt_net.module
          crt_net = crt_net.state_dict()
          crt_net_keys = set(crt_net.keys())
          load_net_keys = set(load_net.keys())
  
          if crt_net_keys != load_net_keys:
              logger.warn('Current net - loaded net:')
              for v in sorted(list(crt_net_keys - load_net_keys)):
                  logger.warn(f'  {v}')
              logger.warn('Loaded net - current net:')
              for v in sorted(list(load_net_keys - crt_net_keys)):
                  logger.warn(f'  {v}')
  
          # check the size for the same keys
          if not strict:
              common_keys = crt_net_keys & load_net_keys
              for k in common_keys:
                  if crt_net[k].size() != load_net[k].size():
                      logger.warn(
                          f'Size different, ignore [{k}]: crt_net: '
                          f'{crt_net[k].shape}; load_net: {load_net[k].shape}')
                      load_net[k + '.ignore'] = load_net.pop(k)
  
      def load_network(self, net, load_path, strict=True):
          """Load network.
  
          Args:
              load_path (str): The path of networks to be loaded.
              net (nn.Module): Network.
              strict (bool): Whether strictly loaded.
          """
          if isinstance(net, nn.DataParallel) or isinstance(
                  net, DistributedDataParallel):
              net = net.module
          net_cls_name = net.__class__.__name__
          logger.info(f'Loading {net_cls_name} model from {load_path}.')
          load_net = torch.load(load_path)
          # remove unnecessary 'module.'
          for k, v in load_net.items():
              if k.startswith('module.'):
                  load_net[k[7:]] = v
                  load_net.pop(k)
          self._print_different_keys_loading(net, load_net, strict)
          net.load_state_dict(load_net, strict=strict)
  
      @master_only
      def save_training_state(self, epoch, current_iter):
          """Save training states during training, which will be used for
          resuming.
  
          Args:
              epoch (int): Current epoch.
              current_iter (int): Current iteration.
          """
          if current_iter != -1:
              state = {
                  'epoch': epoch,
                  'iter': current_iter,
                  'optimizers': [],
                  'schedulers': []
              }
              for o in self.optimizers:
                  state['optimizers'].append(o.state_dict())
              for s in self.schedulers:
                  state['schedulers'].append(s.state_dict())
              save_filename = f'{current_iter}.state'
              save_path = os.path.join(self.opt['path']['training_state'],
                                       save_filename)
              torch.save(state, save_path)
  
      def resume_training(self, resume_state):
          """Reload the optimizers and schedulers for resumed training.
  
          Args:
              resume_state (dict): Resume state.
          """
          resume_optimizers = resume_state['optimizers']
          resume_schedulers = resume_state['schedulers']
          assert len(resume_optimizers) == len(
              self.optimizers), 'Wrong lengths of optimizers'
          assert len(resume_schedulers) == len(
              self.schedulers), 'Wrong lengths of schedulers'
          for i, o in enumerate(resume_optimizers):
              self.optimizers[i].load_state_dict(o)
          for i, s in enumerate(resume_schedulers):
              self.schedulers[i].load_state_dict(s)
  