| | from model import common
|
| | from model import attention
|
| | import torch.nn as nn
|
| |
|
| | def make_model(args, parent=False):
|
| | if args.dilation:
|
| | from model import dilated
|
| | return PAEDSR(args, dilated.dilated_conv)
|
| | else:
|
| | return PAEDSR(args)
|
| |
|
| | class PAEDSR(nn.Module):
|
| | def __init__(self, args, conv=common.default_conv):
|
| | super(PAEDSR, self).__init__()
|
| |
|
| | n_resblock = args.n_resblocks
|
| | n_feats = args.n_feats
|
| | kernel_size = 3
|
| | scale = args.scale[0]
|
| | act = nn.ReLU(True)
|
| |
|
| | rgb_mean = (0.4488, 0.4371, 0.4040)
|
| | rgb_std = (1.0, 1.0, 1.0)
|
| | self.sub_mean = common.MeanShift(args.rgb_range, rgb_mean, rgb_std)
|
| | self.msa = attention.PyramidAttention(channel=256, reduction=8,res_scale=args.res_scale);
|
| |
|
| | m_head = [conv(args.n_colors, n_feats, kernel_size)]
|
| |
|
| |
|
| | m_body = [
|
| | common.ResBlock(
|
| | conv, n_feats, kernel_size, act=act, res_scale=args.res_scale
|
| | ) for _ in range(n_resblock//2)
|
| | ]
|
| | m_body.append(self.msa)
|
| | for _ in range(n_resblock//2):
|
| | m_body.append( common.ResBlock(
|
| | conv, n_feats, kernel_size, act=act, res_scale=args.res_scale
|
| | ))
|
| | m_body.append(conv(n_feats, n_feats, kernel_size))
|
| |
|
| |
|
| | m_tail = [
|
| | common.Upsampler(conv, scale, n_feats, act=False),
|
| | nn.Conv2d(
|
| | n_feats, args.n_colors, kernel_size,
|
| | padding=(kernel_size//2)
|
| | )
|
| | ]
|
| |
|
| | self.add_mean = common.MeanShift(args.rgb_range, rgb_mean, rgb_std, 1)
|
| |
|
| | self.head = nn.Sequential(*m_head)
|
| | self.body = nn.Sequential(*m_body)
|
| | self.tail = nn.Sequential(*m_tail)
|
| |
|
| | def forward(self, x):
|
| | x = self.sub_mean(x)
|
| | x = self.head(x)
|
| |
|
| | res = self.body(x)
|
| | res += x
|
| |
|
| | x = self.tail(res)
|
| | x = self.add_mean(x)
|
| |
|
| | return x
|
| |
|
| | def load_state_dict(self, state_dict, strict=True):
|
| | 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') == -1:
|
| | 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))
|
| |
|
| |
|
