architecture/discriminator.py

# -*- coding: utf-8 -*-
from torch import nn as nn
from torch.nn import functional as F
from torch.nn.utils import spectral_norm
import torch
import functools

class UNetDiscriminatorSN(nn.Module):
    """Defines a U-Net discriminator with spectral normalization (SN)

    It is used in Real-ESRGAN: Training Real-World Blind Super-Resolution with Pure Synthetic Data.

    Arg:
        num_in_ch (int): Channel number of inputs. Default: 3.
        num_feat (int): Channel number of base intermediate features. Default: 64.
        skip_connection (bool): Whether to use skip connections between U-Net. Default: True.
    """

    def __init__(self, num_in_ch, num_feat=64, skip_connection=True):
        super(UNetDiscriminatorSN, self).__init__()
        self.skip_connection = skip_connection
        norm = spectral_norm
        # the first convolution
        self.conv0 = nn.Conv2d(num_in_ch, num_feat, kernel_size=3, stride=1, padding=1)
        # downsample
        self.conv1 = norm(nn.Conv2d(num_feat, num_feat * 2, 4, 2, 1, bias=False))
        self.conv2 = norm(nn.Conv2d(num_feat * 2, num_feat * 4, 4, 2, 1, bias=False))
        self.conv3 = norm(nn.Conv2d(num_feat * 4, num_feat * 8, 4, 2, 1, bias=False))
        # upsample
        self.conv4 = norm(nn.Conv2d(num_feat * 8, num_feat * 4, 3, 1, 1, bias=False))
        self.conv5 = norm(nn.Conv2d(num_feat * 4, num_feat * 2, 3, 1, 1, bias=False))
        self.conv6 = norm(nn.Conv2d(num_feat * 2, num_feat, 3, 1, 1, bias=False))
        # extra convolutions
        self.conv7 = norm(nn.Conv2d(num_feat, num_feat, 3, 1, 1, bias=False))
        self.conv8 = norm(nn.Conv2d(num_feat, num_feat, 3, 1, 1, bias=False))
        self.conv9 = nn.Conv2d(num_feat, 1, 3, 1, 1)

    def forward(self, x):

        # downsample
        x0 = F.leaky_relu(self.conv0(x), negative_slope=0.2, inplace=True)
        x1 = F.leaky_relu(self.conv1(x0), negative_slope=0.2, inplace=True)
        x2 = F.leaky_relu(self.conv2(x1), negative_slope=0.2, inplace=True)
        x3 = F.leaky_relu(self.conv3(x2), negative_slope=0.2, inplace=True)

        # upsample
        x3 = F.interpolate(x3, scale_factor=2, mode='bilinear', align_corners=False)
        x4 = F.leaky_relu(self.conv4(x3), negative_slope=0.2, inplace=True)

        if self.skip_connection:
            x4 = x4 + x2
        x4 = F.interpolate(x4, scale_factor=2, mode='bilinear', align_corners=False)
        x5 = F.leaky_relu(self.conv5(x4), negative_slope=0.2, inplace=True)

        if self.skip_connection:
            x5 = x5 + x1
        x5 = F.interpolate(x5, scale_factor=2, mode='bilinear', align_corners=False)
        x6 = F.leaky_relu(self.conv6(x5), negative_slope=0.2, inplace=True)

        if self.skip_connection:
            x6 = x6 + x0

        # extra convolutions
        out = F.leaky_relu(self.conv7(x6), negative_slope=0.2, inplace=True)
        out = F.leaky_relu(self.conv8(out), negative_slope=0.2, inplace=True)
        out = self.conv9(out)

        return out



def get_conv_layer(input_nc, ndf, kernel_size, stride, padding, bias=True, use_sn=False):
    if not use_sn:
        return nn.Conv2d(input_nc, ndf, kernel_size=kernel_size, stride=stride, padding=padding, bias=bias)
    return spectral_norm(nn.Conv2d(input_nc, ndf, kernel_size=kernel_size, stride=stride, padding=padding, bias=bias))


class PatchDiscriminator(nn.Module):
    """Defines a PatchGAN discriminator, the receptive field of default config is 70x70.

    Args:
        use_sn (bool): Use spectra_norm or not, if use_sn is True, then norm_type should be none.
    """

    def __init__(self,
                 num_in_ch,
                 num_feat=64,
                 num_layers=3,
                 max_nf_mult=8,
                 norm_type='batch',
                 use_sigmoid=False,
                 use_sn=False):
        super(PatchDiscriminator, self).__init__()

        norm_layer = self._get_norm_layer(norm_type)
        if type(norm_layer) == functools.partial:  # no need to use bias as BatchNorm2d has affine parameters
            use_bias = norm_layer.func != nn.BatchNorm2d
        else:
            use_bias = norm_layer != nn.BatchNorm2d

        kw = 4
        padw = 1
        sequence = [
            get_conv_layer(num_in_ch, num_feat, kernel_size=kw, stride=2, padding=padw, use_sn=use_sn),
            nn.LeakyReLU(0.2, True)
        ]
        nf_mult = 1
        nf_mult_prev = 1
        for n in range(1, num_layers):  # gradually increase the number of filters
            nf_mult_prev = nf_mult
            nf_mult = min(2**n, max_nf_mult)
            sequence += [
                get_conv_layer(
                    num_feat * nf_mult_prev,
                    num_feat * nf_mult,
                    kernel_size=kw,
                    stride=2,
                    padding=padw,
                    bias=use_bias,
                    use_sn=use_sn),
                norm_layer(num_feat * nf_mult),
                nn.LeakyReLU(0.2, True)
            ]

        nf_mult_prev = nf_mult
        nf_mult = min(2**num_layers, max_nf_mult)
        sequence += [
            get_conv_layer(
                num_feat * nf_mult_prev,
                num_feat * nf_mult,
                kernel_size=kw,
                stride=1,
                padding=padw,
                bias=use_bias,
                use_sn=use_sn),
            norm_layer(num_feat * nf_mult),
            nn.LeakyReLU(0.2, True)
        ]

        # output 1 channel prediction map 我觉得这个应该就是pixel by pixel的feedback反馈
        sequence += [get_conv_layer(num_feat * nf_mult, 1, kernel_size=kw, stride=1, padding=padw, use_sn=use_sn)]

        if use_sigmoid:
            sequence += [nn.Sigmoid()]
        self.model = nn.Sequential(*sequence)

    def _get_norm_layer(self, norm_type='batch'):
        if norm_type == 'batch':
            norm_layer = functools.partial(nn.BatchNorm2d, affine=True)
        elif norm_type == 'instance':
            norm_layer = functools.partial(nn.InstanceNorm2d, affine=False)
        elif norm_type == 'batchnorm2d':
            norm_layer = nn.BatchNorm2d
        elif norm_type == 'none':
            norm_layer = nn.Identity
        else:
            raise NotImplementedError(f'normalization layer [{norm_type}] is not found')

        return norm_layer

    def forward(self, x):
        return self.model(x)


class MultiScaleDiscriminator(nn.Module):
    """Define a multi-scale discriminator, each discriminator is a instance of PatchDiscriminator.

    Args:
        num_layers (int or list): If the type of this variable is int, then degrade to PatchDiscriminator.
                                  If the type of this variable is list, then the length of the list is
                                  the number of discriminators.
        use_downscale (bool): Progressive downscale the input to feed into different discriminators.
                              If set to True, then the discriminators are usually the same.
    """

    def __init__(self,
                 num_in_ch,
                 num_feat=64,
                 num_layers=[3, 3, 3],
                 max_nf_mult=8,
                 norm_type='none',
                 use_sigmoid=False,
                 use_sn=True,
                 use_downscale=True):
        super(MultiScaleDiscriminator, self).__init__()

        if isinstance(num_layers, int):
            num_layers = [num_layers]

        # check whether the discriminators are the same
        if use_downscale:
            assert len(set(num_layers)) == 1
        self.use_downscale = use_downscale

        self.num_dis = len(num_layers)
        self.dis_list = nn.ModuleList()
        for nl in num_layers:
            self.dis_list.append(
                PatchDiscriminator(
                    num_in_ch,
                    num_feat=num_feat,
                    num_layers=nl,
                    max_nf_mult=max_nf_mult,
                    norm_type=norm_type,
                    use_sigmoid=use_sigmoid,
                    use_sn=use_sn,
                ))

    def forward(self, x):
        outs = []
        h, w = x.size()[2:]

        y = x
        for i in range(self.num_dis):
            if i != 0 and self.use_downscale:
                y = F.interpolate(y, size=(h // 2, w // 2), mode='bilinear', align_corners=True)
                h, w = y.size()[2:]
            outs.append(self.dis_list[i](y))

        return outs


def main():
    from pthflops import count_ops
    from torchsummary import summary
    
    model = UNetDiscriminatorSN(3)
    pytorch_total_params = sum(p.numel() for p in model.parameters())

    # Create a network and a corresponding input
    device = 'cuda'
    inp = torch.rand(1, 3, 400, 400)

    # Count the number of FLOPs
    count_ops(model, inp)
    summary(model.cuda(), (3, 400, 400), batch_size=1)
    # print(f"pathGAN has param {pytorch_total_params//1000} K params")


if __name__ == "__main__":
    main()