Source code for imaginaire.discriminators.multires_patch

# Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# This work is made available under the Nvidia Source Code License-NC.
# To view a copy of this license, check out LICENSE.md
# Copyright (C) 2020 NVIDIA Corporation.  All rights reserved
import functools
import warnings

import numpy as np
import torch
import torch.nn as nn

from imaginaire.layers import Conv2dBlock
from imaginaire.utils.data import (get_paired_input_image_channel_number,
                                   get_paired_input_label_channel_number)
from imaginaire.utils.distributed import master_only_print as print


[docs]class Discriminator(nn.Module): r"""Multi-resolution patch discriminator. Args: dis_cfg (obj): Discriminator definition part of the yaml config file. data_cfg (obj): Data definition part of the yaml config file. """ def __init__(self, dis_cfg, data_cfg): super(Discriminator, self).__init__() print('Multi-resolution patch discriminator initialization.') # We assume the first datum is the ground truth image. image_channels = get_paired_input_image_channel_number(data_cfg) # Calculate number of channels in the input label. num_labels = get_paired_input_label_channel_number(data_cfg) # Build the discriminator. kernel_size = getattr(dis_cfg, 'kernel_size', 3) num_filters = getattr(dis_cfg, 'num_filters', 128) max_num_filters = getattr(dis_cfg, 'max_num_filters', 512) num_discriminators = getattr(dis_cfg, 'num_discriminators', 2) num_layers = getattr(dis_cfg, 'num_layers', 5) activation_norm_type = getattr(dis_cfg, 'activation_norm_type', 'none') weight_norm_type = getattr(dis_cfg, 'weight_norm_type', 'spectral') print('\tBase filter number: %d' % num_filters) print('\tNumber of discriminators: %d' % num_discriminators) print('\tNumber of layers in a discriminator: %d' % num_layers) print('\tWeight norm type: %s' % weight_norm_type) num_input_channels = image_channels + num_labels self.model = MultiResPatchDiscriminator(num_discriminators, kernel_size, num_input_channels, num_filters, num_layers, max_num_filters, activation_norm_type, weight_norm_type) print('Done with the Multi-resolution patch ' 'discriminator initialization.')
[docs] def forward(self, data, net_G_output, real=True): r"""SPADE Generator forward. Args: data (dict): - data (N x C1 x H x W tensor) : Ground truth images. - label (N x C2 x H x W tensor) : Semantic representations. - z (N x style_dims tensor): Gaussian random noise. net_G_output (dict): fake_images (N x C1 x H x W tensor) : Fake images. real (bool): If ``True``, also classifies real images. Otherwise it only classifies generated images to save computation during the generator update. Returns: (tuple): - real_outputs (list): list of output tensors produced by - individual patch discriminators for real images. - real_features (list): list of lists of features produced by individual patch discriminators for real images. - fake_outputs (list): list of output tensors produced by individual patch discriminators for fake images. - fake_features (list): list of lists of features produced by individual patch discriminators for fake images. """ output_x = dict() if 'label' in data: fake_input_x = torch.cat( (data['label'], net_G_output['fake_images']), 1) else: fake_input_x = net_G_output['fake_images'] output_x['fake_outputs'], output_x['fake_features'], _ = \ self.model.forward(fake_input_x) if real: if 'label' in data: real_input_x = torch.cat( (data['label'], data['images']), 1) else: real_input_x = data['images'] output_x['real_outputs'], output_x['real_features'], _ = \ self.model.forward(real_input_x) return output_x
[docs]class MultiResPatchDiscriminator(nn.Module): r"""Multi-resolution patch discriminator. Args: num_discriminators (int): Num. of discriminators (one per scale). kernel_size (int): Convolution kernel size. num_image_channels (int): Num. of channels in the real/fake image. num_filters (int): Num. of base filters in a layer. num_layers (int): Num. of layers for the patch discriminator. max_num_filters (int): Maximum num. of filters in a layer. activation_norm_type (str): batch_norm/instance_norm/none/.... weight_norm_type (str): none/spectral_norm/weight_norm """ def __init__(self, num_discriminators=3, kernel_size=3, num_image_channels=3, num_filters=64, num_layers=4, max_num_filters=512, activation_norm_type='', weight_norm_type='', **kwargs): super().__init__() for key in kwargs: if key != 'type' and key != 'patch_wise': warnings.warn( "Discriminator argument {} is not used".format(key)) self.discriminators = nn.ModuleList() for i in range(num_discriminators): net_discriminator = NLayerPatchDiscriminator( kernel_size, num_image_channels, num_filters, num_layers, max_num_filters, activation_norm_type, weight_norm_type) self.discriminators.append(net_discriminator) print('Done with the Multi-resolution patch ' 'discriminator initialization.')
[docs] def forward(self, input_x): r"""Multi-resolution patch discriminator forward. Args: input_x (tensor) : Input images. Returns: (tuple): - output_list (list): list of output tensors produced by individual patch discriminators. - features_list (list): list of lists of features produced by individual patch discriminators. - input_list (list): list of downsampled input images. """ input_list = [] output_list = [] features_list = [] input_downsampled = input_x for net_discriminator in self.discriminators: input_list.append(input_downsampled) output, features = net_discriminator(input_downsampled) output_list.append(output) features_list.append(features) input_downsampled = nn.functional.interpolate( input_downsampled, scale_factor=0.5, mode='bilinear', align_corners=True, recompute_scale_factor=True) return output_list, features_list, input_list
[docs]class WeightSharedMultiResPatchDiscriminator(nn.Module): r"""Multi-resolution patch discriminator with shared weights. Args: num_discriminators (int): Num. of discriminators (one per scale). kernel_size (int): Convolution kernel size. num_image_channels (int): Num. of channels in the real/fake image. num_filters (int): Num. of base filters in a layer. num_layers (int): Num. of layers for the patch discriminator. max_num_filters (int): Maximum num. of filters in a layer. activation_norm_type (str): batch_norm/instance_norm/none/.... weight_norm_type (str): none/spectral_norm/weight_norm """ def __init__(self, num_discriminators=3, kernel_size=3, num_image_channels=3, num_filters=64, num_layers=4, max_num_filters=512, activation_norm_type='', weight_norm_type='', **kwargs): super().__init__() for key in kwargs: if key != 'type' and key != 'patch_wise': warnings.warn( "Discriminator argument {} is not used".format(key)) self.num_discriminators = num_discriminators self.discriminator = NLayerPatchDiscriminator( kernel_size, num_image_channels, num_filters, num_layers, max_num_filters, activation_norm_type, weight_norm_type) print('Done with the Weight-Shared Multi-resolution patch ' 'discriminator initialization.')
[docs] def forward(self, input_x): r"""Multi-resolution patch discriminator forward. Args: input_x (tensor) : Input images. Returns: (tuple): - output_list (list): list of output tensors produced by individual patch discriminators. - features_list (list): list of lists of features produced by individual patch discriminators. - input_list (list): list of downsampled input images. """ input_list = [] output_list = [] features_list = [] input_downsampled = input_x for i in range(self.num_discriminators): input_list.append(input_downsampled) output, features = self.discriminator(input_downsampled) output_list.append(output) features_list.append(features) input_downsampled = nn.functional.interpolate( input_downsampled, scale_factor=0.5, mode='bilinear', align_corners=True) return output_list, features_list, input_list
[docs]class NLayerPatchDiscriminator(nn.Module): r"""Patch Discriminator constructor. Args: kernel_size (int): Convolution kernel size. num_input_channels (int): Num. of channels in the real/fake image. num_filters (int): Num. of base filters in a layer. num_layers (int): Num. of layers for the patch discriminator. max_num_filters (int): Maximum num. of filters in a layer. activation_norm_type (str): batch_norm/instance_norm/none/.... weight_norm_type (str): none/spectral_norm/weight_norm """ def __init__(self, kernel_size, num_input_channels, num_filters, num_layers, max_num_filters, activation_norm_type, weight_norm_type): super(NLayerPatchDiscriminator, self).__init__() self.num_layers = num_layers padding = int(np.floor((kernel_size - 1.0) / 2)) nonlinearity = 'leakyrelu' base_conv2d_block = \ functools.partial(Conv2dBlock, kernel_size=kernel_size, padding=padding, weight_norm_type=weight_norm_type, activation_norm_type=activation_norm_type, nonlinearity=nonlinearity, # inplace_nonlinearity=True, order='CNA') layers = [[base_conv2d_block( num_input_channels, num_filters, stride=2)]] for n in range(num_layers): num_filters_prev = num_filters num_filters = min(num_filters * 2, max_num_filters) stride = 2 if n < (num_layers - 1) else 1 layers += [[base_conv2d_block(num_filters_prev, num_filters, stride=stride)]] layers += [[Conv2dBlock(num_filters, 1, 3, 1, padding, weight_norm_type=weight_norm_type)]] for n in range(len(layers)): setattr(self, 'layer' + str(n), nn.Sequential(*layers[n]))
[docs] def forward(self, input_x): r"""Patch Discriminator forward. Args: input_x (N x C x H1 x W2 tensor): Concatenation of images and semantic representations. Returns: (tuple): - output (N x 1 x H2 x W2 tensor): Discriminator output value. Before the sigmoid when using NSGAN. - features (list): lists of tensors of the intermediate activations. """ res = [input_x] for n in range(self.num_layers + 2): layer = getattr(self, 'layer' + str(n)) x = res[-1] res.append(layer(x)) output = res[-1] features = res[1:-1] return output, features