CoolFace
Apppublic

MLBench/ReaLens

sourceHugging Facemitupdated 8mo agoView on Hugging Face
0likes
networks.py589 linesDownload Raw Back to models
1import torch2import torch.nn as nn3from torch.nn import init4import functools5from torch.optim import lr_scheduler6 7 8###############################################################################9# Helper Functions10###############################################################################11 12 13class Identity(nn.Module):14    def forward(self, x):15        return x16 17 18def get_norm_layer(norm_type="instance"):19    """Return a normalization layer20 21    Parameters:22        norm_type (str) -- the name of the normalization layer: batch | instance | none23 24    For BatchNorm, we use learnable affine parameters and track running statistics (mean/stddev).25    For InstanceNorm, we do not use learnable affine parameters. We do not track running statistics.26    """27    if norm_type == "batch":28        norm_layer = functools.partial(nn.BatchNorm2d, affine=True, track_running_stats=True)29    elif norm_type == "syncbatch":30        norm_layer = functools.partial(nn.SyncBatchNorm, affine=True, track_running_stats=True)31    elif norm_type == "instance":32        norm_layer = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False)33    elif norm_type == "none":34 35        def norm_layer(x):36            return Identity()37 38    else:39        raise NotImplementedError("normalization layer [%s] is not found" % norm_type)40    return norm_layer41 42 43def get_scheduler(optimizer, opt):44    """Return a learning rate scheduler45 46    Parameters:47        optimizer          -- the optimizer of the network48        opt (option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions.49                              opt.lr_policy is the name of learning rate policy: linear | step | plateau | cosine50 51    For 'linear', we keep the same learning rate for the first <opt.n_epochs> epochs52    and linearly decay the rate to zero over the next <opt.n_epochs_decay> epochs.53    For other schedulers (step, plateau, and cosine), we use the default PyTorch schedulers.54    See https://pytorch.org/docs/stable/optim.html for more details.55    """56    if opt.lr_policy == "linear":57 58        def lambda_rule(epoch):59            lr_l = 1.0 - max(0, epoch + opt.epoch_count - opt.n_epochs) / float(opt.n_epochs_decay + 1)60            return lr_l61 62        scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule)63    elif opt.lr_policy == "step":64        scheduler = lr_scheduler.StepLR(optimizer, step_size=opt.lr_decay_iters, gamma=0.1)65    elif opt.lr_policy == "plateau":66        scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode="min", factor=0.2, threshold=0.01, patience=5)67    elif opt.lr_policy == "cosine":68        scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=opt.n_epochs, eta_min=0)69    else:70        return NotImplementedError("learning rate policy [%s] is not implemented", opt.lr_policy)71    return scheduler72 73 74def init_weights(net, init_type="normal", init_gain=0.02):75    """Initialize network weights.76 77    Parameters:78        net (network)   -- network to be initialized79        init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal80        init_gain (float)    -- scaling factor for normal, xavier and orthogonal.81 82    We use 'normal' in the original pix2pix and CycleGAN paper. But xavier and kaiming might83    work better for some applications. Feel free to try yourself.84    """85 86    def init_func(m):  # define the initialization function87        classname = m.__class__.__name__88        if hasattr(m, "weight") and (classname.find("Conv") != -1 or classname.find("Linear") != -1):89            if init_type == "normal":90                init.normal_(m.weight.data, 0.0, init_gain)91            elif init_type == "xavier":92                init.xavier_normal_(m.weight.data, gain=init_gain)93            elif init_type == "kaiming":94                init.kaiming_normal_(m.weight.data, a=0, mode="fan_in")95            elif init_type == "orthogonal":96                init.orthogonal_(m.weight.data, gain=init_gain)97            else:98                raise NotImplementedError("initialization method [%s] is not implemented" % init_type)99            if hasattr(m, "bias") and m.bias is not None:100                init.constant_(m.bias.data, 0.0)101        elif classname.find("BatchNorm2d") != -1:  # BatchNorm Layer's weight is not a matrix; only normal distribution applies.102            init.normal_(m.weight.data, 1.0, init_gain)103            init.constant_(m.bias.data, 0.0)104 105    print("initialize network with %s" % init_type)106    net.apply(init_func)  # apply the initialization function <init_func>107 108 109def init_net(net, init_type="normal", init_gain=0.02):110    """Initialize a network: 1. register CPU/GPU device; 2. initialize the network weights111    Parameters:112        net (network)      -- the network to be initialized113        init_type (str)    -- the name of an initialization method: normal | xavier | kaiming | orthogonal114        gain (float)       -- scaling factor for normal, xavier and orthogonal.115 116    Return an initialized network.117    """118    import os119 120    if torch.cuda.is_available():121        if "LOCAL_RANK" in os.environ:122            local_rank = int(os.environ["LOCAL_RANK"])123            net.to(local_rank)124            print(f"Initialized with device cuda:{local_rank}")125        else:126            net.to(0)127            print("Initialized with device cuda:0")128    init_weights(net, init_type, init_gain=init_gain)129    return net130 131 132def define_G(input_nc, output_nc, ngf, netG, norm="batch", use_dropout=False, init_type="normal", init_gain=0.02):133    """Create a generator134 135    Parameters:136        input_nc (int) -- the number of channels in input images137        output_nc (int) -- the number of channels in output images138        ngf (int) -- the number of filters in the last conv layer139        netG (str) -- the architecture's name: resnet_9blocks | resnet_6blocks | unet_128 | unet_256140        norm (str) -- the name of normalization layers used in the network: batch | instance | none141        use_dropout (bool) -- if use dropout layers.142        init_type (str)    -- the name of our initialization method.143        init_gain (float)  -- scaling factor for normal, xavier and orthogonal.144 145    Returns a generator146    """147    net = None148    norm_layer = get_norm_layer(norm_type=norm)149 150    if netG == "resnet_9blocks":151        net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=9)152    elif netG == "resnet_6blocks":153        net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, n_blocks=6)154    elif netG == "unet_128":155        net = UnetGenerator(input_nc, output_nc, 7, ngf, norm_layer=norm_layer, use_dropout=use_dropout)156    elif netG == "unet_256":157        net = UnetGenerator(input_nc, output_nc, 8, ngf, norm_layer=norm_layer, use_dropout=use_dropout)158    else:159        raise NotImplementedError("Generator model name [%s] is not recognized" % netG)160    return net161 162 163def define_D(input_nc, ndf, netD, n_layers_D=3, norm="batch", init_type="normal", init_gain=0.02):164    """Create a discriminator165 166    Parameters:167        input_nc (int)     -- the number of channels in input images168        ndf (int)          -- the number of filters in the first conv layer169        netD (str)         -- the architecture's name: basic | n_layers | pixel170        n_layers_D (int)   -- the number of conv layers in the discriminator; effective when netD=='n_layers'171        norm (str)         -- the type of normalization layers used in the network.172        init_type (str)    -- the name of the initialization method.173        init_gain (float)  -- scaling factor for normal, xavier and orthogonal.174 175    Returns a discriminator176 177    Our current implementation provides three types of discriminators:178        [basic]: 'PatchGAN' classifier described in the original pix2pix paper.179        It can classify whether 70×70 overlapping patches are real or fake.180        Such a patch-level discriminator architecture has fewer parameters181        than a full-image discriminator and can work on arbitrarily-sized images182        in a fully convolutional fashion.183 184        [n_layers]: With this mode, you can specify the number of conv layers in the discriminator185        with the parameter <n_layers_D> (default=3 as used in [basic] (PatchGAN).)186 187        [pixel]: 1x1 PixelGAN discriminator can classify whether a pixel is real or not.188        It encourages greater color diversity but has no effect on spatial statistics.189 190    The discriminator has been initialized by <init_net>. It uses Leakly RELU for non-linearity.191    """192    net = None193    norm_layer = get_norm_layer(norm_type=norm)194 195    if netD == "basic":  # default PatchGAN classifier196        net = NLayerDiscriminator(input_nc, ndf, n_layers=3, norm_layer=norm_layer)197    elif netD == "n_layers":  # more options198        net = NLayerDiscriminator(input_nc, ndf, n_layers_D, norm_layer=norm_layer)199    elif netD == "pixel":  # classify if each pixel is real or fake200        net = PixelDiscriminator(input_nc, ndf, norm_layer=norm_layer)201    else:202        raise NotImplementedError("Discriminator model name [%s] is not recognized" % netD)203    return net204 205 206##############################################################################207# Classes208##############################################################################209class GANLoss(nn.Module):210    """Define different GAN objectives.211 212    The GANLoss class abstracts away the need to create the target label tensor213    that has the same size as the input.214    """215 216    def __init__(self, gan_mode, target_real_label=1.0, target_fake_label=0.0):217        """Initialize the GANLoss class.218 219        Parameters:220            gan_mode (str) - - the type of GAN objective. It currently supports vanilla, lsgan, and wgangp.221            target_real_label (bool) - - label for a real image222            target_fake_label (bool) - - label of a fake image223 224        Note: Do not use sigmoid as the last layer of Discriminator.225        LSGAN needs no sigmoid. vanilla GANs will handle it with BCEWithLogitsLoss.226        """227        super(GANLoss, self).__init__()228        self.register_buffer("real_label", torch.tensor(target_real_label))229        self.register_buffer("fake_label", torch.tensor(target_fake_label))230        self.gan_mode = gan_mode231        if gan_mode == "lsgan":232            self.loss = nn.MSELoss()233        elif gan_mode == "vanilla":234            self.loss = nn.BCEWithLogitsLoss()235        elif gan_mode in ["wgangp"]:236            self.loss = None237        else:238            raise NotImplementedError("gan mode %s not implemented" % gan_mode)239 240    def get_target_tensor(self, prediction, target_is_real):241        """Create label tensors with the same size as the input.242 243        Parameters:244            prediction (tensor) - - tpyically the prediction from a discriminator245            target_is_real (bool) - - if the ground truth label is for real images or fake images246 247        Returns:248            A label tensor filled with ground truth label, and with the size of the input249        """250 251        if target_is_real:252            target_tensor = self.real_label253        else:254            target_tensor = self.fake_label255        return target_tensor.expand_as(prediction)256 257    def __call__(self, prediction, target_is_real):258        """Calculate loss given Discriminator's output and grount truth labels.259 260        Parameters:261            prediction (tensor) - - tpyically the prediction output from a discriminator262            target_is_real (bool) - - if the ground truth label is for real images or fake images263 264        Returns:265            the calculated loss.266        """267        if self.gan_mode in ["lsgan", "vanilla"]:268            target_tensor = self.get_target_tensor(prediction, target_is_real)269            loss = self.loss(prediction, target_tensor)270        elif self.gan_mode == "wgangp":271            if target_is_real:272                loss = -prediction.mean()273            else:274                loss = prediction.mean()275        return loss276 277 278def cal_gradient_penalty(netD, real_data, fake_data, device, type="mixed", constant=1.0, lambda_gp=10.0):279    """Calculate the gradient penalty loss, used in WGAN-GP paper https://arxiv.org/abs/1704.00028280 281    Arguments:282        netD (network)              -- discriminator network283        real_data (tensor array)    -- real images284        fake_data (tensor array)    -- generated images from the generator285        device (str)                -- GPU / CPU286        type (str)                  -- if we mix real and fake data or not [real | fake | mixed].287        constant (float)            -- the constant used in formula ( ||gradient||_2 - constant)^2288        lambda_gp (float)           -- weight for this loss289 290    Returns the gradient penalty loss291    """292    if lambda_gp > 0.0:293        if type == "real":  # either use real images, fake images, or a linear interpolation of two.294            interpolatesv = real_data295        elif type == "fake":296            interpolatesv = fake_data297        elif type == "mixed":298            alpha = torch.rand(real_data.shape[0], 1, device=device)299            alpha = alpha.expand(real_data.shape[0], real_data.nelement() // real_data.shape[0]).contiguous().view(*real_data.shape)300            interpolatesv = alpha * real_data + ((1 - alpha) * fake_data)301        else:302            raise NotImplementedError(f"{type} not implemented")303        interpolatesv.requires_grad_(True)304        disc_interpolates = netD(interpolatesv)305        gradients = torch.autograd.grad(outputs=disc_interpolates, inputs=interpolatesv, grad_outputs=torch.ones(disc_interpolates.size()).to(device), create_graph=True, retain_graph=True, only_inputs=True)306        gradients = gradients[0].view(real_data.size(0), -1)  # flat the data307        gradient_penalty = (((gradients + 1e-16).norm(2, dim=1) - constant) ** 2).mean() * lambda_gp  # added eps308        return gradient_penalty, gradients309    else:310        return 0.0, None311 312 313class ResnetGenerator(nn.Module):314    """Resnet-based generator that consists of Resnet blocks between a few downsampling/upsampling operations.315 316    We adapt Torch code and idea from Justin Johnson's neural style transfer project(https://github.com/jcjohnson/fast-neural-style)317    """318 319    def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=6, padding_type="reflect"):320        """Construct a Resnet-based generator321 322        Parameters:323            input_nc (int)      -- the number of channels in input images324            output_nc (int)     -- the number of channels in output images325            ngf (int)           -- the number of filters in the last conv layer326            norm_layer          -- normalization layer327            use_dropout (bool)  -- if use dropout layers328            n_blocks (int)      -- the number of ResNet blocks329            padding_type (str)  -- the name of padding layer in conv layers: reflect | replicate | zero330        """331        assert n_blocks >= 0332        super(ResnetGenerator, self).__init__()333        if type(norm_layer) == functools.partial:334            use_bias = norm_layer.func == nn.InstanceNorm2d335        else:336            use_bias = norm_layer == nn.InstanceNorm2d337 338        model = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias), norm_layer(ngf), nn.ReLU(True)]339 340        n_downsampling = 2341        for i in range(n_downsampling):  # add downsampling layers342            mult = 2**i343            model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1, bias=use_bias), norm_layer(ngf * mult * 2), nn.ReLU(True)]344 345        mult = 2**n_downsampling346        for i in range(n_blocks):  # add ResNet blocks347 348            model += [ResnetBlock(ngf * mult, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)]349 350        for i in range(n_downsampling):  # add upsampling layers351            mult = 2 ** (n_downsampling - i)352            model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2), kernel_size=3, stride=2, padding=1, output_padding=1, bias=use_bias), norm_layer(int(ngf * mult / 2)), nn.ReLU(True)]353        model += [nn.ReflectionPad2d(3)]354        model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)]355        model += [nn.Tanh()]356 357        self.model = nn.Sequential(*model)358 359    def forward(self, input):360        """Standard forward"""361        return self.model(input)362 363 364class ResnetBlock(nn.Module):365    """Define a Resnet block"""366 367    def __init__(self, dim, padding_type, norm_layer, use_dropout, use_bias):368        """Initialize the Resnet block369 370        A resnet block is a conv block with skip connections371        We construct a conv block with build_conv_block function,372        and implement skip connections in <forward> function.373        Original Resnet paper: https://arxiv.org/pdf/1512.03385.pdf374        """375        super(ResnetBlock, self).__init__()376        self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, use_dropout, use_bias)377 378    def build_conv_block(self, dim, padding_type, norm_layer, use_dropout, use_bias):379        """Construct a convolutional block.380 381        Parameters:382            dim (int)           -- the number of channels in the conv layer.383            padding_type (str)  -- the name of padding layer: reflect | replicate | zero384            norm_layer          -- normalization layer385            use_dropout (bool)  -- if use dropout layers.386            use_bias (bool)     -- if the conv layer uses bias or not387 388        Returns a conv block (with a conv layer, a normalization layer, and a non-linearity layer (ReLU))389        """390        conv_block = []391        p = 0392        if padding_type == "reflect":393            conv_block += [nn.ReflectionPad2d(1)]394        elif padding_type == "replicate":395            conv_block += [nn.ReplicationPad2d(1)]396        elif padding_type == "zero":397            p = 1398        else:399            raise NotImplementedError("padding [%s] is not implemented" % padding_type)400 401        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim), nn.ReLU(True)]402        if use_dropout:403            conv_block += [nn.Dropout(0.5)]404 405        p = 0406        if padding_type == "reflect":407            conv_block += [nn.ReflectionPad2d(1)]408        elif padding_type == "replicate":409            conv_block += [nn.ReplicationPad2d(1)]410        elif padding_type == "zero":411            p = 1412        else:413            raise NotImplementedError("padding [%s] is not implemented" % padding_type)414        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim)]415 416        return nn.Sequential(*conv_block)417 418    def forward(self, x):419        """Forward function (with skip connections)"""420        out = x + self.conv_block(x)  # add skip connections421        return out422 423 424class UnetGenerator(nn.Module):425    """Create a Unet-based generator"""426 427    def __init__(self, input_nc, output_nc, num_downs, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False):428        """Construct a Unet generator429        Parameters:430            input_nc (int)  -- the number of channels in input images431            output_nc (int) -- the number of channels in output images432            num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7,433                                image of size 128x128 will become of size 1x1 # at the bottleneck434            ngf (int)       -- the number of filters in the last conv layer435            norm_layer      -- normalization layer436 437        We construct the U-Net from the innermost layer to the outermost layer.438        It is a recursive process.439        """440        super(UnetGenerator, self).__init__()441        # construct unet structure442        unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True)  # add the innermost layer443        for i in range(num_downs - 5):  # add intermediate layers with ngf * 8 filters444            unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer, use_dropout=use_dropout)445        # gradually reduce the number of filters from ngf * 8 to ngf446        unet_block = UnetSkipConnectionBlock(ngf * 4, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer)447        unet_block = UnetSkipConnectionBlock(ngf * 2, ngf * 4, input_nc=None, submodule=unet_block, norm_layer=norm_layer)448        unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer)449        self.model = UnetSkipConnectionBlock(output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer)  # add the outermost layer450 451    def forward(self, input):452        """Standard forward"""453        return self.model(input)454 455 456class UnetSkipConnectionBlock(nn.Module):457    """Defines the Unet submodule with skip connection.458    X -------------------identity----------------------459    |-- downsampling -- |submodule| -- upsampling --|460    """461 462    def __init__(self, outer_nc, inner_nc, input_nc=None, submodule=None, outermost=False, innermost=False, norm_layer=nn.BatchNorm2d, use_dropout=False):463        """Construct a Unet submodule with skip connections.464 465        Parameters:466            outer_nc (int) -- the number of filters in the outer conv layer467            inner_nc (int) -- the number of filters in the inner conv layer468            input_nc (int) -- the number of channels in input images/features469            submodule (UnetSkipConnectionBlock) -- previously defined submodules470            outermost (bool)    -- if this module is the outermost module471            innermost (bool)    -- if this module is the innermost module472            norm_layer          -- normalization layer473            use_dropout (bool)  -- if use dropout layers.474        """475        super(UnetSkipConnectionBlock, self).__init__()476        self.outermost = outermost477        if type(norm_layer) == functools.partial:478            use_bias = norm_layer.func == nn.InstanceNorm2d479        else:480            use_bias = norm_layer == nn.InstanceNorm2d481        if input_nc is None:482            input_nc = outer_nc483        downconv = nn.Conv2d(input_nc, inner_nc, kernel_size=4, stride=2, padding=1, bias=use_bias)484        downrelu = nn.LeakyReLU(0.2, True)485        downnorm = norm_layer(inner_nc)486        uprelu = nn.ReLU(True)487        upnorm = norm_layer(outer_nc)488 489        if outermost:490            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc, kernel_size=4, stride=2, padding=1)491            down = [downconv]492            up = [uprelu, upconv, nn.Tanh()]493            model = down + [submodule] + up494        elif innermost:495            upconv = nn.ConvTranspose2d(inner_nc, outer_nc, kernel_size=4, stride=2, padding=1, bias=use_bias)496            down = [downrelu, downconv]497            up = [uprelu, upconv, upnorm]498            model = down + up499        else:500            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc, kernel_size=4, stride=2, padding=1, bias=use_bias)501            down = [downrelu, downconv, downnorm]502            up = [uprelu, upconv, upnorm]503 504            if use_dropout:505                model = down + [submodule] + up + [nn.Dropout(0.5)]506            else:507                model = down + [submodule] + up508 509        self.model = nn.Sequential(*model)510 511    def forward(self, x):512        if self.outermost:513            return self.model(x)514        else:  # add skip connections515            return torch.cat([x, self.model(x)], 1)516 517 518class NLayerDiscriminator(nn.Module):519    """Defines a PatchGAN discriminator"""520 521    def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d):522        """Construct a PatchGAN discriminator523 524        Parameters:525            input_nc (int)  -- the number of channels in input images526            ndf (int)       -- the number of filters in the last conv layer527            n_layers (int)  -- the number of conv layers in the discriminator528            norm_layer      -- normalization layer529        """530        super(NLayerDiscriminator, self).__init__()531        if type(norm_layer) == functools.partial:  # no need to use bias as BatchNorm2d has affine parameters532            use_bias = norm_layer.func == nn.InstanceNorm2d533        else:534            use_bias = norm_layer == nn.InstanceNorm2d535 536        kw = 4537        padw = 1538        sequence = [nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)]539        nf_mult = 1540        nf_mult_prev = 1541        for n in range(1, n_layers):  # gradually increase the number of filters542            nf_mult_prev = nf_mult543            nf_mult = min(2**n, 8)544            sequence += [nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=2, padding=padw, bias=use_bias), norm_layer(ndf * nf_mult), nn.LeakyReLU(0.2, True)]545 546        nf_mult_prev = nf_mult547        nf_mult = min(2**n_layers, 8)548        sequence += [nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=1, padding=padw, bias=use_bias), norm_layer(ndf * nf_mult), nn.LeakyReLU(0.2, True)]549 550        sequence += [nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)]  # output 1 channel prediction map551        self.model = nn.Sequential(*sequence)552 553    def forward(self, input):554        """Standard forward."""555        return self.model(input)556 557 558class PixelDiscriminator(nn.Module):559    """Defines a 1x1 PatchGAN discriminator (pixelGAN)"""560 561    def __init__(self, input_nc, ndf=64, norm_layer=nn.BatchNorm2d):562        """Construct a 1x1 PatchGAN discriminator563 564        Parameters:565            input_nc (int)  -- the number of channels in input images566            ndf (int)       -- the number of filters in the last conv layer567            norm_layer      -- normalization layer568        """569        super(PixelDiscriminator, self).__init__()570        if type(norm_layer) == functools.partial:  # no need to use bias as BatchNorm2d has affine parameters571            use_bias = norm_layer.func == nn.InstanceNorm2d572        else:573            use_bias = norm_layer == nn.InstanceNorm2d574 575        self.net = [576            nn.Conv2d(input_nc, ndf, kernel_size=1, stride=1, padding=0),577            nn.LeakyReLU(0.2, True),578            nn.Conv2d(ndf, ndf * 2, kernel_size=1, stride=1, padding=0, bias=use_bias),579            norm_layer(ndf * 2),580            nn.LeakyReLU(0.2, True),581            nn.Conv2d(ndf * 2, 1, kernel_size=1, stride=1, padding=0, bias=use_bias),582        ]583 584        self.net = nn.Sequential(*self.net)585 586    def forward(self, input):587        """Standard forward."""588        return self.net(input)589