| from model import common
|
| import torch.nn as nn
|
| from lambda_networks import LambdaLayer
|
| import torch.cuda.amp as amp
|
|
|
| def make_model(args, parent=False):
|
| return LAMBDANET(args)
|
|
|
| class MCheck(nn.Module):
|
| def __init__(self):
|
| super(MCheck, self).__init__()
|
| self.kill = False
|
| self.shape = None
|
|
|
| def forward(self, x):
|
| if self.kill:
|
| raise NotImplementedError("Shape checked here", self.shape, x.shape)
|
| self.kill = True
|
| self.shape = x.shape
|
| return x
|
|
|
| class LAMBDANET(nn.Module):
|
| def __init__(self, args, conv=common.default_conv):
|
| super(LAMBDANET, 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)
|
| layer = LambdaLayer(
|
| dim = n_feats,
|
| dim_out = n_feats,
|
| r = 23,
|
| dim_k = 16,
|
| heads = 4,
|
| dim_u = 4
|
| )
|
|
|
|
|
| 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(layer)
|
| 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)
|
|
|
| self.recurrence = args.recurrence
|
| self.detach = args.detach
|
|
|
| self.amp = args.amp
|
|
|
| def forward(self, x):
|
|
|
| with amp.autocast(self.amp):
|
| last_output=x
|
| output_lst=[None]*self.recurrence
|
| for i in range(self.recurrence):
|
| if self.detach:
|
| last_output=last_output.detach()
|
| out = self.head(last_output)
|
| res = self.body(out)
|
| res += out
|
| output = self.tail(res) + last_output
|
| output_lst[i]=output
|
| last_output=output
|
|
|
| return output_lst
|
|
|
| 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))
|
|
|
|
|
