MLBench/ReaLens
0
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 