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