|
|
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))
|
|
|
|
|
|
|
