|
from torch import nn as nn
|
|
|
|
from basicsr.utils.registry import ARCH_REGISTRY
|
|
|
|
|
|
@ARCH_REGISTRY.register()
|
|
class VGGStyleDiscriminator128(nn.Module):
|
|
"""VGG style discriminator with input size 128 x 128.
|
|
|
|
It is used to train SRGAN and ESRGAN.
|
|
|
|
Args:
|
|
num_in_ch (int): Channel number of inputs. Default: 3.
|
|
num_feat (int): Channel number of base intermediate features.
|
|
Default: 64.
|
|
"""
|
|
|
|
def __init__(self, num_in_ch, num_feat):
|
|
super(VGGStyleDiscriminator128, self).__init__()
|
|
|
|
self.conv0_0 = nn.Conv2d(num_in_ch, num_feat, 3, 1, 1, bias=True)
|
|
self.conv0_1 = nn.Conv2d(num_feat, num_feat, 4, 2, 1, bias=False)
|
|
self.bn0_1 = nn.BatchNorm2d(num_feat, affine=True)
|
|
|
|
self.conv1_0 = nn.Conv2d(num_feat, num_feat * 2, 3, 1, 1, bias=False)
|
|
self.bn1_0 = nn.BatchNorm2d(num_feat * 2, affine=True)
|
|
self.conv1_1 = nn.Conv2d(num_feat * 2, num_feat * 2, 4, 2, 1, bias=False)
|
|
self.bn1_1 = nn.BatchNorm2d(num_feat * 2, affine=True)
|
|
|
|
self.conv2_0 = nn.Conv2d(num_feat * 2, num_feat * 4, 3, 1, 1, bias=False)
|
|
self.bn2_0 = nn.BatchNorm2d(num_feat * 4, affine=True)
|
|
self.conv2_1 = nn.Conv2d(num_feat * 4, num_feat * 4, 4, 2, 1, bias=False)
|
|
self.bn2_1 = nn.BatchNorm2d(num_feat * 4, affine=True)
|
|
|
|
self.conv3_0 = nn.Conv2d(num_feat * 4, num_feat * 8, 3, 1, 1, bias=False)
|
|
self.bn3_0 = nn.BatchNorm2d(num_feat * 8, affine=True)
|
|
self.conv3_1 = nn.Conv2d(num_feat * 8, num_feat * 8, 4, 2, 1, bias=False)
|
|
self.bn3_1 = nn.BatchNorm2d(num_feat * 8, affine=True)
|
|
|
|
self.conv4_0 = nn.Conv2d(num_feat * 8, num_feat * 8, 3, 1, 1, bias=False)
|
|
self.bn4_0 = nn.BatchNorm2d(num_feat * 8, affine=True)
|
|
self.conv4_1 = nn.Conv2d(num_feat * 8, num_feat * 8, 4, 2, 1, bias=False)
|
|
self.bn4_1 = nn.BatchNorm2d(num_feat * 8, affine=True)
|
|
|
|
self.linear1 = nn.Linear(num_feat * 8 * 4 * 4, 100)
|
|
self.linear2 = nn.Linear(100, 1)
|
|
|
|
|
|
self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
|
|
|
|
def forward(self, x):
|
|
assert x.size(2) == 128 and x.size(3) == 128, (f'Input spatial size must be 128x128, '
|
|
f'but received {x.size()}.')
|
|
|
|
feat = self.lrelu(self.conv0_0(x))
|
|
feat = self.lrelu(self.bn0_1(self.conv0_1(feat)))
|
|
|
|
feat = self.lrelu(self.bn1_0(self.conv1_0(feat)))
|
|
feat = self.lrelu(self.bn1_1(self.conv1_1(feat)))
|
|
|
|
feat = self.lrelu(self.bn2_0(self.conv2_0(feat)))
|
|
feat = self.lrelu(self.bn2_1(self.conv2_1(feat)))
|
|
|
|
feat = self.lrelu(self.bn3_0(self.conv3_0(feat)))
|
|
feat = self.lrelu(self.bn3_1(self.conv3_1(feat)))
|
|
|
|
feat = self.lrelu(self.bn4_0(self.conv4_0(feat)))
|
|
feat = self.lrelu(self.bn4_1(self.conv4_1(feat)))
|
|
|
|
feat = feat.view(feat.size(0), -1)
|
|
feat = self.lrelu(self.linear1(feat))
|
|
out = self.linear2(feat)
|
|
return out
|
|
|
