File size: 1,464 Bytes
add457a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 |
# Copyright (C) 2020 NVIDIA Corporation. 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
from torch import nn
def get_nonlinearity_layer(nonlinearity_type, inplace):
r"""Return a nonlinearity layer.
Args:
nonlinearity_type (str):
Type of nonlinear activation function.
``'none'``, ``'relu'``, ``'leakyrelu'``, ``'prelu'``,
``'tanh'`` , ``'sigmoid'`` or ``'softmax'``.
inplace (bool): If ``True``, set ``inplace=True`` when initializing
the nonlinearity layer.
"""
if nonlinearity_type == 'relu':
nonlinearity = nn.ReLU(inplace=inplace)
elif nonlinearity_type == 'leakyrelu':
nonlinearity = nn.LeakyReLU(0.2, inplace=inplace)
elif nonlinearity_type == 'prelu':
nonlinearity = nn.PReLU()
elif nonlinearity_type == 'tanh':
nonlinearity = nn.Tanh()
elif nonlinearity_type == 'sigmoid':
nonlinearity = nn.Sigmoid()
elif nonlinearity_type.startswith('softmax'):
dim = nonlinearity_type.split(',')[1] if ',' in nonlinearity_type else 1
nonlinearity = nn.Softmax(dim=int(dim))
elif nonlinearity_type == 'none' or nonlinearity_type == '':
nonlinearity = None
else:
raise ValueError('Nonlinearity %s is not recognized' %
nonlinearity_type)
return nonlinearity
|