Source code for imaginaire.discriminators.mlp_multiclass

# 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
import functools

import numpy as np
import torch.nn as nn

from imaginaire.layers import LinearBlock

[docs]class Discriminator(nn.Module): r"""Multi-layer Perceptron Classifier constructor. 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__() num_input_channels = dis_cfg.input_dims num_labels = dis_cfg.num_labels num_layers = getattr(dis_cfg, 'num_layers', 5) num_filters = getattr(dis_cfg, 'num_filters', 512) activation_norm_type = getattr(dis_cfg, 'activation_norm_type', 'batch_norm') nonlinearity = getattr(dis_cfg, 'nonlinearity', 'leakyrelu') base_linear_block = \ functools.partial(LinearBlock, activation_norm_type=activation_norm_type, nonlinearity=nonlinearity, order='CNA') dropout_ratio = 0.1 layers = [base_linear_block(num_input_channels, num_filters), nn.Dropout(dropout_ratio)] for n in range(num_layers): dropout_ratio *= 1.5 dropout_ratio = np.min([dropout_ratio, 0.5]) layers += [base_linear_block(num_filters, num_filters), nn.Dropout(dropout_ratio)] layers += [LinearBlock(num_filters, num_labels)] self.model = nn.Sequential(*layers)
[docs] def forward(self, data): r"""Patch Discriminator forward. Args: data (dict): - data (N x -1 tensor): We will reshape the tensor to this format. Returns: (dict): - results (N x C tensor): Output scores before softmax. """ input_x = data['data'] bs = input_x.size()[0] input_x = input_x.view(bs, -1) pre_softmax_scores = self.model(input_x) outputs = dict() outputs['results'] = pre_softmax_scores return outputs