|
|
import torch |
|
|
import torch.nn as nn |
|
|
|
|
|
from . import common |
|
|
from .ResNet import ResNet |
|
|
|
|
|
|
|
|
def build_model(args): |
|
|
return MSResNet(args) |
|
|
|
|
|
class conv_end(nn.Module): |
|
|
def __init__(self, in_channels=3, out_channels=3, kernel_size=5, ratio=2): |
|
|
super(conv_end, self).__init__() |
|
|
|
|
|
modules = [ |
|
|
common.default_conv(in_channels, out_channels, kernel_size), |
|
|
nn.PixelShuffle(ratio) |
|
|
] |
|
|
|
|
|
self.uppath = nn.Sequential(*modules) |
|
|
|
|
|
def forward(self, x): |
|
|
return self.uppath(x) |
|
|
|
|
|
class MSResNet(nn.Module): |
|
|
def __init__(self, args): |
|
|
super(MSResNet, self).__init__() |
|
|
|
|
|
self.rgb_range = args.rgb_range |
|
|
self.mean = self.rgb_range / 2 |
|
|
|
|
|
self.n_resblocks = args.n_resblocks |
|
|
self.n_feats = args.n_feats |
|
|
self.kernel_size = args.kernel_size |
|
|
|
|
|
self.n_scales = args.n_scales |
|
|
|
|
|
self.body_models = nn.ModuleList([ |
|
|
ResNet(args, 3, 3, mean_shift=False), |
|
|
]) |
|
|
for _ in range(1, self.n_scales): |
|
|
self.body_models.insert(0, ResNet(args, 6, 3, mean_shift=False)) |
|
|
|
|
|
self.conv_end_models = nn.ModuleList([None]) |
|
|
for _ in range(1, self.n_scales): |
|
|
self.conv_end_models += [conv_end(3, 12)] |
|
|
|
|
|
def forward(self, input_pyramid): |
|
|
|
|
|
scales = range(self.n_scales-1, -1, -1) |
|
|
|
|
|
for s in scales: |
|
|
input_pyramid[s] = input_pyramid[s] - self.mean |
|
|
|
|
|
output_pyramid = [None] * self.n_scales |
|
|
|
|
|
input_s = input_pyramid[-1] |
|
|
for s in scales: |
|
|
output_pyramid[s] = self.body_models[s](input_s) |
|
|
if s > 0: |
|
|
up_feat = self.conv_end_models[s](output_pyramid[s]) |
|
|
input_s = torch.cat((input_pyramid[s-1], up_feat), 1) |
|
|
|
|
|
for s in scales: |
|
|
output_pyramid[s] = output_pyramid[s] + self.mean |
|
|
|
|
|
return output_pyramid |
|
|
|
