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