text: import math
  
  import torch
  import torch.nn as nn
  from torch.nn.parameter import Parameter
  def make_model(args, parent=False):
      return ARAN(args)
  
  class ChannelAttention(nn.Module):
      def __init__(self, in_planes, ratio=16):
          super(ChannelAttention, self).__init__()
          self.avg_pool = nn.AdaptiveAvgPool2d(1)
          self.max_pool = nn.AdaptiveMaxPool2d(1)
  
          self.fc1   = nn.Conv2d(in_planes, in_planes // 16, 1, bias=False)
          self.relu1 = nn.ReLU()
          self.fc2   = nn.Conv2d(in_planes // 16, in_planes, 1, bias=False)
  
          self.sigmoid = nn.Sigmoid()
  
      def forward(self, x):
          avg_out = self.fc2(self.relu1(self.fc1(self.avg_pool(x))))
          max_out = self.fc2(self.relu1(self.fc1(self.max_pool(x))))
          out = avg_out + max_out
          return self.sigmoid(out) * x
  
  class SpatialAttention(nn.Module):
      def __init__(self, kernel_size=7):
          super(SpatialAttention, self).__init__()
  
          assert kernel_size in (3, 7), 'kernel size must be 3 or 7'
          padding = 3 if kernel_size == 7 else 1
  
          self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
          self.sigmoid = nn.Sigmoid()
  
      def forward(self, x):
          avg_out = torch.mean(x, dim=1, keepdim=True)
          max_out, _ = torch.max(x, dim=1, keepdim=True)
          y = torch.cat([avg_out, max_out], dim=1)
          y = self.conv1(y)
          return self.sigmoid(y) * x
  
  class Block(nn.Module):
      def __init__(
          self, n_feats, kernel_size, block_feats, wn, res_scale=1, act=nn.ReLU(True)):
          super(Block, self).__init__()
          self.res_scale = 1
          
          body = []
          body.append(
              wn(nn.Conv2d(n_feats, block_feats, 1, padding=0)))
          body.append(act)
          body.append(
              wn(nn.Conv2d(block_feats, block_feats, kernel_size, padding=kernel_size//2, groups=64)))
          body.append(act)
          body.append(
              wn(nn.Conv2d(block_feats, n_feats, 1, padding=0)))
  
          body.append(ChannelAttention(n_feats, 16))
          body.append(SpatialAttention())
  
          self.body = nn.Sequential(*body)
  
      def forward(self, x):
          res = self.body(x) * self.res_scale
          res += x
          return res
  
  class ARAN(nn.Module):
      def __init__(self, args):
          super(ARAN, self).__init__()
          # hyper-params
          self.args = args
          scale = args.scale[0]
          n_resblocks = 16
          n_feats = 128
          block_feats = 256
          kernel_size = 3
          act = nn.ReLU(True)
          
          # wn = lambda x: x
          wn = lambda x: torch.nn.utils.weight_norm(x)
  
          self.rgb_mean = torch.autograd.Variable(torch.FloatTensor(
              [args.r_mean, args.g_mean, args.b_mean])).view([1, 3, 1, 1])
  
          # define head module
          head = []
          head.append(
              wn(nn.Conv2d(args.n_colors, n_feats, 3, padding=3//2)))
  
          # define body module
          body = []
          for i in range(n_resblocks):
              body.append(
                  Block(n_feats, kernel_size, block_feats, wn=wn, res_scale=args.res_scale, act=act))
  
          # define tail module
          tail = []
          out_feats = scale*scale*args.n_colors
          tail.append(
              wn(nn.Conv2d(n_feats, out_feats, 3, padding=3//2)))
          tail.append(nn.PixelShuffle(scale))
  
          skip = []
          skip.append(
              wn(nn.Conv2d(args.n_colors, out_feats, 5, padding=5//2))
          )
          skip.append(nn.PixelShuffle(scale))
  
          # make object members
          self.head = nn.Sequential(*head)
          self.body = nn.Sequential(*body)
          self.tail = nn.Sequential(*tail)
          self.skip = nn.Sequential(*skip)
  
      def forward(self, x):
          x = (x - self.rgb_mean.cuda()*255)/127.5
          s = self.skip(x)
          x = self.head(x)
          x = self.body(x)
          x = self.tail(x)
          x += s
          x = x*127.5 + self.rgb_mean.cuda()*255
          return x
      
      def load_state_dict(self, state_dict, strict=False):
          own_state = self.state_dict()
          for name, param in state_dict.items():
              if name in own_state:
                  if isinstance(param, nn.Parameter):
                      param = param.data
                  try:
                      own_state[name].copy_(param)
                  except Exception:
                      if name.find('tail') or name.find('skip') >= 0:
                          print('Replace pre-trained upsampler to new one...')
                      else:
                          raise RuntimeError('While copying the parameter named {}, '
                                             'whose dimensions in the model are {} and '
                                             'whose dimensions in the checkpoint are {}.'
                                             .format(name, own_state[name].size(), param.size()))
              elif strict:
                  if name.find('tail') == -1:
                      raise KeyError('unexpected key "{}" in state_dict'
                                     .format(name))
  
          if strict:
              missing = set(own_state.keys()) - set(state_dict.keys())
              if len(missing) > 0:
                  raise KeyError('missing keys in state_dict: "{}"'.format(missing))
  