forked from mrq/DL-Art-School
92 lines
3.0 KiB
Python
92 lines
3.0 KiB
Python
|
from models.archs.panet import common
|
||
|
|
from models.archs.panet import attention
|
||
|
|
import torch.nn as nn
|
||
|
|
from utils.util import checkpoint
|
||
|
|
|
||
|
|
|
||
|
|
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)
|
||
|
|
self.msa = attention.PyramidAttention()
|
||
|
|
# define head module
|
||
|
|
m_head = [conv(args.n_colors, n_feats, kernel_size)]
|
||
|
|
|
||
|
|
# define body module
|
||
|
|
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(self.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))
|
||
|
|
|
||
|
|
# define tail module
|
||
|
|
# m_tail = [
|
||
|
|
# common.Upsampler(conv, scale, n_feats, act=False),
|
||
|
|
# conv(n_feats, args.n_colors, kernel_size)
|
||
|
|
# ]
|
||
|
|
m_tail = [
|
||
|
|
common.Upsampler(conv, scale, n_feats, act=False),
|
||
|
|
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.ModuleList(m_body)
|
||
|
|
self.tail = nn.Sequential(*m_tail)
|
||
|
|
|
||
|
|
def forward(self, x):
|
||
|
|
# x = self.sub_mean(x)
|
||
|
|
x = self.head(x)
|
||
|
|
|
||
|
|
res = x
|
||
|
|
for b in self.body:
|
||
|
|
if b == self.msa:
|
||
|
|
if __name__ == '__main__':
|
||
|
|
res = self.msa(res)
|
||
|
|
else:
|
||
|
|
res = checkpoint(b, res)
|
||
|
|
|
||
|
|
res += x
|
||
|
|
|
||
|
|
x = checkpoint(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))
|