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