| from model import common
|
| from model import attention
|
| import torch.nn as nn
|
|
|
| def make_model(args, parent=False):
|
| return PANET(args)
|
|
|
| class PANET(nn.Module):
|
| def __init__(self, args, conv=common.default_conv):
|
| super(PANET, self).__init__()
|
|
|
| n_resblocks = args.n_resblocks
|
| n_feats = args.n_feats
|
| kernel_size = 3
|
| scale = args.scale[0]
|
|
|
| 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)
|
| msa = attention.PyramidAttention()
|
|
|
| m_head = [conv(args.n_colors, n_feats, kernel_size)]
|
|
|
|
|
| m_body = [
|
| common.ResBlock(
|
| conv, n_feats, kernel_size, nn.PReLU(), res_scale=args.res_scale
|
| ) for _ in range(n_resblocks//2)
|
| ]
|
| m_body.append(msa)
|
| for i in range(n_resblocks//2):
|
| m_body.append(common.ResBlock(conv,n_feats,kernel_size,nn.PReLU(),res_scale=args.res_scale))
|
|
|
| m_body.append(conv(n_feats, n_feats, kernel_size))
|
|
|
|
|
|
|
|
|
|
|
|
|
| m_tail = [
|
| conv(n_feats, args.n_colors, kernel_size)
|
| ]
|
|
|
| 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.head(x)
|
|
|
| res = self.body(x)
|
|
|
| res += x
|
|
|
| x = self.tail(res)
|
|
|
|
|
| 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))
|
|
|
|
|
