ai-art

 1# https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/models/networks.py
 2
 3import functools
 4import torch.nn as nn
 5
 6from core.taming.modules.discriminator import ActNorm
 7
 8
 9class NLayerDiscriminator(nn.Module):
10    """Defines a PatchGAN discriminator as in Pix2Pix"""
11    def __init__(self, input_nc=3, ndf=64, n_layers=3, use_actnorm=False):
12        """Construct a PatchGAN discriminator
13        Parameters:
14            input_nc (int)  -- the number of channels in input images
15            ndf (int)       -- the number of filters in the last conv layer
16            n_layers (int)  -- the number of conv layers in the discriminator
17            norm_layer      -- normalization layer
18        """
19        super(NLayerDiscriminator, self).__init__()
20        if not use_actnorm:
21            norm_layer = nn.BatchNorm2d
22        else:
23            norm_layer = ActNorm
24        if type(norm_layer) == functools.partial:  # no need to use bias as BatchNorm2d has affine parameters
25            use_bias = norm_layer.func != nn.BatchNorm2d
26        else:
27            use_bias = norm_layer != nn.BatchNorm2d
28
29        kw = 4
30        padw = 1
31        sequence = [nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)]
32        nf_mult = 1
33        nf_mult_prev = 1
34        for n in range(1, n_layers):  # gradually increase the number of filters
35            nf_mult_prev = nf_mult
36            nf_mult = min(2 ** n, 8)
37            sequence += [
38                nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=2, padding=padw, bias=use_bias),
39                norm_layer(ndf * nf_mult),
40                nn.LeakyReLU(0.2, True)
41            ]
42
43        nf_mult_prev = nf_mult
44        nf_mult = min(2 ** n_layers, 8)
45        sequence += [
46            nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=1, padding=padw, bias=use_bias),
47            norm_layer(ndf * nf_mult),
48            nn.LeakyReLU(0.2, True)
49        ]
50
51        sequence += [
52            nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)]  # output 1 channel prediction map
53        self.main = nn.Sequential(*sequence)
54
55    def forward(self, input):
56        """Standard forward."""
57        return self.main(input)