|
import torch.nn as nn |
|
import torch |
|
|
|
class DenseResidualBlock(nn.Module): |
|
""" |
|
The core module of paper: (Residual Dense Network for Image Super-Resolution, CVPR 18) |
|
""" |
|
|
|
def __init__(self, filters, res_scale=0.2): |
|
super(DenseResidualBlock, self).__init__() |
|
self.res_scale = res_scale |
|
|
|
def block(in_features, non_linearity=True): |
|
layers = [nn.Conv2d(in_features, filters, 3, 1, 1, bias=True)] |
|
if non_linearity: |
|
layers += [nn.LeakyReLU()] |
|
return nn.Sequential(*layers) |
|
|
|
self.b1 = block(in_features=1 * filters) |
|
self.b2 = block(in_features=2 * filters) |
|
self.b3 = block(in_features=3 * filters) |
|
self.b4 = block(in_features=4 * filters) |
|
self.b5 = block(in_features=5 * filters, non_linearity=False) |
|
self.blocks = [self.b1, self.b2, self.b3, self.b4, self.b5] |
|
|
|
def forward(self, x): |
|
inputs = x |
|
for block in self.blocks: |
|
out = block(inputs) |
|
inputs = torch.cat([inputs, out], 1) |
|
return out.mul(self.res_scale) + x |
|
|
|
|
|
class ResidualInResidualDenseBlock(nn.Module): |
|
def __init__(self, filters, res_scale=0.2): |
|
super(ResidualInResidualDenseBlock, self).__init__() |
|
self.res_scale = res_scale |
|
self.dense_blocks = nn.Sequential( |
|
DenseResidualBlock(filters), DenseResidualBlock(filters), DenseResidualBlock(filters) |
|
) |
|
|
|
def forward(self, x): |
|
return self.dense_blocks(x).mul(self.res_scale) + x |
|
|
|
|
|
class GeneratorRRDB(nn.Module): |
|
def __init__(self, channels, filters=64, num_res_blocks=16, num_upsample=2): |
|
super(GeneratorRRDB, self).__init__() |
|
|
|
|
|
self.conv1 = nn.Conv2d(channels, filters, kernel_size=3, stride=1, padding=1) |
|
|
|
self.res_blocks = nn.Sequential(*[ResidualInResidualDenseBlock(filters) for _ in range(num_res_blocks)]) |
|
|
|
self.conv2 = nn.Conv2d(filters, filters, kernel_size=3, stride=1, padding=1) |
|
|
|
upsample_layers = [] |
|
for _ in range(num_upsample): |
|
upsample_layers += [ |
|
nn.Conv2d(filters, filters * 4, kernel_size=3, stride=1, padding=1), |
|
nn.LeakyReLU(), |
|
nn.PixelShuffle(upscale_factor=2), |
|
] |
|
self.upsampling = nn.Sequential(*upsample_layers) |
|
|
|
self.conv3 = nn.Sequential( |
|
nn.Conv2d(filters, filters, kernel_size=3, stride=1, padding=1), |
|
nn.LeakyReLU(), |
|
nn.Conv2d(filters, channels, kernel_size=3, stride=1, padding=1), |
|
) |
|
|
|
def forward(self, x): |
|
out1 = self.conv1(x) |
|
out = self.res_blocks(out1) |
|
out2 = self.conv2(out) |
|
out = torch.add(out1, out2) |
|
out = self.upsampling(out) |
|
out = self.conv3(out) |
|
return out |