text: import os
  from importlib import import_module
  
  import torch
  import torch.nn as nn
  import torch.nn.parallel as P
  import torch.utils.model_zoo
  
  class Model(nn.Module):
      def __init__(self, args, ckp):
          super(Model, self).__init__()
          print('Making model...')
  
          self.scale = args.scale
          self.idx_scale = 0
          self.input_large = (args.model == 'VDSR')
          self.self_ensemble = args.self_ensemble
          self.chop = args.chop
          self.precision = args.precision
          self.cpu = args.cpu
          self.device = torch.device('cpu' if args.cpu else 'cuda')
          self.n_GPUs = args.n_GPUs
          self.save_models = args.save_models
  
          module = import_module('model.' + args.model.lower())
          self.model = module.make_model(args).to(self.device)
          if args.precision == 'half':
              self.model.half()
  
          self.load(
              ckp.get_path('model'),
              pre_train=args.pre_train,
              resume=args.resume,
              cpu=args.cpu
          )
          print(self.model, file=ckp.log_file)
  
      def forward(self, x, idx_scale):
          self.idx_scale = idx_scale
          if hasattr(self.model, 'set_scale'):
              self.model.set_scale(idx_scale)
  
          if self.training:
              if self.n_GPUs > 1:
                  return P.data_parallel(self.model, x, range(self.n_GPUs))
              else:
                  return self.model(x)
          else:
              if self.chop:
                  forward_function = self.forward_chop
              else:
                  forward_function = self.model.forward
  
              if self.self_ensemble:
                  return self.forward_x8(x, forward_function=forward_function)
              else:
                  return forward_function(x)
  
      def save(self, apath, epoch, is_best=False):
          save_dirs = [os.path.join(apath, 'model_latest.pt')]
  
          if is_best:
              save_dirs.append(os.path.join(apath, 'model_best.pt'))
          if self.save_models:
              save_dirs.append(
                  os.path.join(apath, 'model_{}.pt'.format(epoch))
              )
  
          for s in save_dirs:
              torch.save(self.model.state_dict(), s)
  
      def load(self, apath, pre_train='', resume=-1, cpu=False):
          load_from = None
          kwargs = {}
          if cpu:
              kwargs = {'map_location': lambda storage, loc: storage}
  
          if resume == -1:
              load_from = torch.load(
                  os.path.join(apath, 'model_latest.pt'),
                  **kwargs
              )
          elif resume == 0:
              if pre_train == 'download':
                  print('Download the model')
                  dir_model = os.path.join('..', 'models')
                  os.makedirs(dir_model, exist_ok=True)
                  load_from = torch.utils.model_zoo.load_url(
                      self.model.url,
                      model_dir=dir_model,
                      **kwargs
                  )
              elif pre_train:
                  print('Load the model from {}'.format(pre_train))
                  load_from = torch.load(pre_train, **kwargs)
          else:
              load_from = torch.load(
                  os.path.join(apath, 'model_{}.pt'.format(resume)),
                  **kwargs
              )
  
          if load_from:
              self.model.load_state_dict(load_from, strict=False)
  
      # def forward_chop(self, *args, shave=10, min_size=160000):
      #     scale = 1 if self.input_large else self.scale[self.idx_scale]
      #     n_GPUs = min(self.n_GPUs, 4)
      #     # height, width
      #     h, w = args[0].size()[-2:]
      #
      #     top = slice(0, h//2 + shave)
      #     bottom = slice(h - h//2 - shave, h)
      #     left = slice(0, w//2 + shave)
      #     right = slice(w - w//2 - shave, w)
      #     x_chops = [torch.cat([
      #         a[..., top, left],
      #         a[..., top, right],
      #         a[..., bottom, left],
      #         a[..., bottom, right]
      #     ]) for a in args]
      #
      #     y_chops = []
      #     if h * w < 4 * min_size:
      #         for i in range(0, 4, n_GPUs):
      #             x = [x_chop[i:(i + n_GPUs)] for x_chop in x_chops]
      #             y = P.data_parallel(self.model, *x, range(n_GPUs))
      #             if not isinstance(y, list): y = [y]
      #             if not y_chops:
      #                 y_chops = [[c for c in _y.chunk(n_GPUs, dim=0)] for _y in y]
      #             else:
      #                 for y_chop, _y in zip(y_chops, y):
      #                     y_chop.extend(_y.chunk(n_GPUs, dim=0))
      #     else:
      #         for p in zip(*x_chops):
      #             y = self.forward_chop(*p, shave=shave, min_size=min_size)
      #             if not isinstance(y, list): y = [y]
      #             if not y_chops:
      #                 y_chops = [[_y] for _y in y]
      #             else:
      #                 for y_chop, _y in zip(y_chops, y): y_chop.append(_y)
      #
      #     h *= scale
      #     w *= scale
      #     top = slice(0, h//2)
      #     bottom = slice(h - h//2, h)
      #     bottom_r = slice(h//2 - h, None)
      #     left = slice(0, w//2)
      #     right = slice(w - w//2, w)
      #     right_r = slice(w//2 - w, None)
      #
      #     # batch size, number of color channels
      #     b, c = y_chops[0][0].size()[:-2]
      #     y = [y_chop[0].new(b, c, h, w) for y_chop in y_chops]
      #     for y_chop, _y in zip(y_chops, y):
      #         _y[..., top, left] = y_chop[0][..., top, left]
      #         _y[..., top, right] = y_chop[1][..., top, right_r]
      #         _y[..., bottom, left] = y_chop[2][..., bottom_r, left]
      #         _y[..., bottom, right] = y_chop[3][..., bottom_r, right_r]
      #
      #     if len(y) == 1: y = y[0]
      #
      #     return y
      #
      # def forward_x8(self, *args, forward_function=None):
      #     def _transform(v, op):
      #         if self.precision != 'single': v = v.float()
      #
      #         v2np = v.data.cpu().numpy()
      #         if op == 'v':
      #             tfnp = v2np[:, :, :, ::-1].copy()
      #         elif op == 'h':
      #             tfnp = v2np[:, :, ::-1, :].copy()
      #         elif op == 't':
      #             tfnp = v2np.transpose((0, 1, 3, 2)).copy()
      #
      #         ret = torch.Tensor(tfnp).to(self.device)
      #         if self.precision == 'half': ret = ret.half()
      #
      #         return ret
      #
      #     list_x = []
      #     for a in args:
      #         x = [a]
      #         for tf in 'v', 'h', 't': x.extend([_transform(_x, tf) for _x in x])
      #
      #         list_x.append(x)
      #
      #     list_y = []
      #     for x in zip(*list_x):
      #         y = forward_function(*x)
      #         if not isinstance(y, list): y = [y]
      #         if not list_y:
      #             list_y = [[_y] for _y in y]
      #         else:
      #             for _list_y, _y in zip(list_y, y): _list_y.append(_y)
      #
      #     for _list_y in list_y:
      #         for i in range(len(_list_y)):
      #             if i > 3:
      #                 _list_y[i] = _transform(_list_y[i], 't')
      #             if i % 4 > 1:
      #                 _list_y[i] = _transform(_list_y[i], 'h')
      #             if (i % 4) % 2 == 1:
      #                 _list_y[i] = _transform(_list_y[i], 'v')
      #
      #     y = [torch.cat(_y, dim=0).mean(dim=0, keepdim=True) for _y in list_y]
      #     if len(y) == 1: y = y[0]
      #
      #     return y
  
      #
      # def forward_chop(self, x, shave=10, min_size=160000):
      #     scale = self.scale[self.idx_scale]
      #     n_GPUs = min(self.n_GPUs, 4)
      #     b, c, h, w = x.size()
      #     h_half, w_half = h // 2, w // 2
      #     h_size, w_size = h_half + shave, w_half + shave
      #     lr_list = [
      #         x[:, :, 0:h_size, 0:w_size],
      #         x[:, :, 0:h_size, (w - w_size):w],
      #         x[:, :, (h - h_size):h, 0:w_size],
      #         x[:, :, (h - h_size):h, (w - w_size):w]]
      #
      #     if w_size * h_size < min_size:
      #         sr_list = []
      #         for i in range(0, 4, n_GPUs):
      #             lr_batch = torch.cat(lr_list[i:(i + n_GPUs)], dim=0)
      #             sr_batch = self.model(lr_batch)
      #             sr_list.extend(sr_batch.chunk(n_GPUs, dim=0))
      #     else:
      #         sr_list = [
      #             self.forward_chop(patch, shave=shave, min_size=min_size) \
      #             for patch in lr_list
      #         ]
      #
      #     h, w = scale * h, scale * w
      #     h_half, w_half = scale * h_half, scale * w_half
      #     h_size, w_size = scale * h_size, scale * w_size
      #     shave *= scale
      #
      #     output = x.new(b, c, h, w)
      #     output[:, :, 0:h_half, 0:w_half] \
      #         = sr_list[0][:, :, 0:h_half, 0:w_half]
      #     output[:, :, 0:h_half, w_half:w] \
      #         = sr_list[1][:, :, 0:h_half, (w_size - w + w_half):w_size]
      #     output[:, :, h_half:h, 0:w_half] \
      #         = sr_list[2][:, :, (h_size - h + h_half):h_size, 0:w_half]
      #     output[:, :, h_half:h, w_half:w] \
      #         = sr_list[3][:, :, (h_size - h + h_half):h_size, (w_size - w + w_half):w_size]
      #
      #     return output
      #
      # def forward_x8(self, x, forward_function):
      #     def _transform(v, op):
      #         if self.precision != 'single': v = v.float()
      #
      #         v2np = v.data.cpu().numpy()
      #         if op == 'v':
      #             tfnp = v2np[:, :, :, ::-1].copy()
      #         elif op == 'h':
      #             tfnp = v2np[:, :, ::-1, :].copy()
      #         elif op == 't':
      #             tfnp = v2np.transpose((0, 1, 3, 2)).copy()
      #
      #         ret = torch.Tensor(tfnp).to(self.device)
      #         if self.precision == 'half': ret = ret.half()
      #
      #         return ret
      #
      #     lr_list = [x]
      #     for tf in 'v', 'h', 't':
      #         lr_list.extend([_transform(t, tf) for t in lr_list])
      #
      #     sr_list = [forward_function(aug) for aug in lr_list]
      #     for i in range(len(sr_list)):
      #         if i > 3:
      #             sr_list[i] = _transform(sr_list[i], 't')
      #         if i % 4 > 1:
      #             sr_list[i] = _transform(sr_list[i], 'h')
      #         if (i % 4) % 2 == 1:
      #             sr_list[i] = _transform(sr_list[i], 'v')
      #
      #     output_cat = torch.cat(sr_list, dim=0)
      #     output = output_cat.mean(dim=0, keepdim=True)
      #
      #     return output
      def forward_chop(self, x, shave=10, min_size=160000):
          scale = self.scale[self.idx_scale]
          n_GPUs = min(self.n_GPUs, 4)
          b, c, h, w = x.size()
          h_half, w_half = h // 2, w // 2
          h_size, w_size = h_half + shave, w_half + shave
          lr_list = [
              x[:, :, 0:h_size, 0:w_size],
              x[:, :, 0:h_size, (w - w_size):w],
              x[:, :, (h - h_size):h, 0:w_size],
              x[:, :, (h - h_size):h, (w - w_size):w]]
  
          if w_size * h_size < min_size:
              sr_list = []
              for i in range(0, 4, n_GPUs):
                  lr_batch = torch.cat(lr_list[i:(i + n_GPUs)], dim=0)
                  sr_batch = self.model(lr_batch)
                  sr_list.extend(sr_batch.chunk(n_GPUs, dim=0))
          else:
              sr_list = [
                  self.forward_chop(patch, shave=shave, min_size=min_size) \
                  for patch in lr_list
              ]
  
          h, w = scale * h, scale * w
          h_half, w_half = scale * h_half, scale * w_half
          h_size, w_size = scale * h_size, scale * w_size
          shave *= scale
  
          output = x.new(b, c, h, w)
          output[:, :, 0:h_half, 0:w_half] \
              = sr_list[0][:, :, 0:h_half, 0:w_half]
          output[:, :, 0:h_half, w_half:w] \
              = sr_list[1][:, :, 0:h_half, (w_size - w + w_half):w_size]
          output[:, :, h_half:h, 0:w_half] \
              = sr_list[2][:, :, (h_size - h + h_half):h_size, 0:w_half]
          output[:, :, h_half:h, w_half:w] \
              = sr_list[3][:, :, (h_size - h + h_half):h_size, (w_size - w + w_half):w_size]
  
          return output
  
      def forward_x8(self, x, forward_function):
          def _transform(v, op):
              if self.precision != 'single': v = v.float()
  
              v2np = v.data.cpu().numpy()
              if op == 'v':
                  tfnp = v2np[:, :, :, ::-1].copy()
              elif op == 'h':
                  tfnp = v2np[:, :, ::-1, :].copy()
              elif op == 't':
                  tfnp = v2np.transpose((0, 1, 3, 2)).copy()
  
              ret = torch.Tensor(tfnp).to(self.device)
              if self.precision == 'half': ret = ret.half()
  
              return ret
  
          lr_list = [x]
          for tf in 'v', 'h', 't':
              lr_list.extend([_transform(t, tf) for t in lr_list])
  
          sr_list = [forward_function(aug) for aug in lr_list]
          for i in range(len(sr_list)):
              if i > 3:
                  sr_list[i] = _transform(sr_list[i], 't')
              if i % 4 > 1:
                  sr_list[i] = _transform(sr_list[i], 'h')
              if (i % 4) % 2 == 1:
                  sr_list[i] = _transform(sr_list[i], 'v')
  
          output_cat = torch.cat(sr_list, dim=0)
          output = output_cat.mean(dim=0, keepdim=True)
  
          return output
  