|
|
|
|
|
|
| from model import common
|
|
|
| import torch
|
| import torch.nn as nn
|
|
|
|
|
| def make_model(args, parent=False):
|
| return DDBPN(args)
|
|
|
| def projection_conv(in_channels, out_channels, scale, up=True):
|
| kernel_size, stride, padding = {
|
| 2: (6, 2, 2),
|
| 4: (8, 4, 2),
|
| 8: (12, 8, 2)
|
| }[scale]
|
| if up:
|
| conv_f = nn.ConvTranspose2d
|
| else:
|
| conv_f = nn.Conv2d
|
|
|
| return conv_f(
|
| in_channels, out_channels, kernel_size,
|
| stride=stride, padding=padding
|
| )
|
|
|
| class DenseProjection(nn.Module):
|
| def __init__(self, in_channels, nr, scale, up=True, bottleneck=True):
|
| super(DenseProjection, self).__init__()
|
| if bottleneck:
|
| self.bottleneck = nn.Sequential(*[
|
| nn.Conv2d(in_channels, nr, 1),
|
| nn.PReLU(nr)
|
| ])
|
| inter_channels = nr
|
| else:
|
| self.bottleneck = None
|
| inter_channels = in_channels
|
|
|
| self.conv_1 = nn.Sequential(*[
|
| projection_conv(inter_channels, nr, scale, up),
|
| nn.PReLU(nr)
|
| ])
|
| self.conv_2 = nn.Sequential(*[
|
| projection_conv(nr, inter_channels, scale, not up),
|
| nn.PReLU(inter_channels)
|
| ])
|
| self.conv_3 = nn.Sequential(*[
|
| projection_conv(inter_channels, nr, scale, up),
|
| nn.PReLU(nr)
|
| ])
|
|
|
| def forward(self, x):
|
| if self.bottleneck is not None:
|
| x = self.bottleneck(x)
|
|
|
| a_0 = self.conv_1(x)
|
| b_0 = self.conv_2(a_0)
|
| e = b_0.sub(x)
|
| a_1 = self.conv_3(e)
|
|
|
| out = a_0.add(a_1)
|
|
|
| return out
|
|
|
| class DDBPN(nn.Module):
|
| def __init__(self, args):
|
| super(DDBPN, self).__init__()
|
| scale = args.scale[0]
|
|
|
| n0 = 128
|
| nr = 32
|
| self.depth = 6
|
|
|
| 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)
|
| initial = [
|
| nn.Conv2d(args.n_colors, n0, 3, padding=1),
|
| nn.PReLU(n0),
|
| nn.Conv2d(n0, nr, 1),
|
| nn.PReLU(nr)
|
| ]
|
| self.initial = nn.Sequential(*initial)
|
|
|
| self.upmodules = nn.ModuleList()
|
| self.downmodules = nn.ModuleList()
|
| channels = nr
|
| for i in range(self.depth):
|
| self.upmodules.append(
|
| DenseProjection(channels, nr, scale, True, i > 1)
|
| )
|
| if i != 0:
|
| channels += nr
|
|
|
| channels = nr
|
| for i in range(self.depth - 1):
|
| self.downmodules.append(
|
| DenseProjection(channels, nr, scale, False, i != 0)
|
| )
|
| channels += nr
|
|
|
| reconstruction = [
|
| nn.Conv2d(self.depth * nr, args.n_colors, 3, padding=1)
|
| ]
|
| self.reconstruction = nn.Sequential(*reconstruction)
|
|
|
| self.add_mean = common.MeanShift(args.rgb_range, rgb_mean, rgb_std, 1)
|
|
|
| def forward(self, x):
|
| x = self.sub_mean(x)
|
| x = self.initial(x)
|
|
|
| h_list = []
|
| l_list = []
|
| for i in range(self.depth - 1):
|
| if i == 0:
|
| l = x
|
| else:
|
| l = torch.cat(l_list, dim=1)
|
| h_list.append(self.upmodules[i](l))
|
| l_list.append(self.downmodules[i](torch.cat(h_list, dim=1)))
|
|
|
| h_list.append(self.upmodules[-1](torch.cat(l_list, dim=1)))
|
| out = self.reconstruction(torch.cat(h_list, dim=1))
|
| out = self.add_mean(out)
|
|
|
| return out
|
|
|
|
|
