|
|
import torch
|
|
|
import torch.nn as nn
|
|
|
|
|
|
from . import common
|
|
|
from .LamResNet import ResNet
|
|
|
|
|
|
|
|
|
def build_model(args):
|
|
|
return RecLamResNet(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 RecLamResNet(nn.Module):
|
|
|
def __init__(self, args):
|
|
|
super(RecLamResNet, self).__init__()
|
|
|
|
|
|
self.rgb_range = args.rgb_range
|
|
|
self.mean = self.rgb_range / 2
|
|
|
self.is_detach=args.detach
|
|
|
|
|
|
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_model = ResNet(args, 3, 3, mean_shift=False)
|
|
|
|
|
|
def forward(self, input_lst):
|
|
|
|
|
|
input_lst[0] = input_lst[0] - self.mean
|
|
|
output_lst = [None] * self.n_scales
|
|
|
last_output = input_lst[0]
|
|
|
for i in range(self.n_scales):
|
|
|
if self.is_detach:
|
|
|
last_output=last_output.detach()
|
|
|
output = self.body_model(last_output) + last_output
|
|
|
output_lst[self.n_scales-i-1] = output + self.mean
|
|
|
last_output = output
|
|
|
return output_lst
|
|
|
|
