Source code for ding.torch_utils.network.activation
import math
from collections.abc import Callable
import torch
import torch.nn as nn
[docs]class Lambda(nn.Module):
"""
Overview:
A custom lambda module for constructing custom layers.
Interfaces:
``__init__``, ``forward``.
"""
[docs] def __init__(self, f: Callable):
"""
Overview:
Initialize the lambda module with a given function.
Arguments:
- f (:obj:`Callable`): a python function
"""
super(Lambda, self).__init__()
self.f = f
[docs] def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Overview:
Compute the function of the input tensor.
Arguments:
- x (:obj:`torch.Tensor`): The input tensor.
"""
return self.f(x)
[docs]class GLU(nn.Module):
"""
Overview:
Gating Linear Unit (GLU), a specific type of activation function, which is first proposed in
[Language Modeling with Gated Convolutional Networks](https://arxiv.org/pdf/1612.08083.pdf).
Interfaces:
``__init__``, ``forward``.
"""
[docs] def __init__(self, input_dim: int, output_dim: int, context_dim: int, input_type: str = 'fc') -> None:
"""
Overview:
Initialize the GLU module.
Arguments:
- input_dim (:obj:`int`): The dimension of the input tensor.
- output_dim (:obj:`int`): The dimension of the output tensor.
- context_dim (:obj:`int`): The dimension of the context tensor.
- input_type (:obj:`str`): The type of input, now supports ['fc', 'conv2d']
"""
super(GLU, self).__init__()
assert (input_type in ['fc', 'conv2d'])
if input_type == 'fc':
self.layer1 = nn.Linear(context_dim, input_dim)
self.layer2 = nn.Linear(input_dim, output_dim)
elif input_type == 'conv2d':
self.layer1 = nn.Conv2d(context_dim, input_dim, 1, 1, 0)
self.layer2 = nn.Conv2d(input_dim, output_dim, 1, 1, 0)
[docs] def forward(self, x: torch.Tensor, context: torch.Tensor) -> torch.Tensor:
"""
Overview:
Compute the GLU transformation of the input tensor.
Arguments:
- x (:obj:`torch.Tensor`): The input tensor.
- context (:obj:`torch.Tensor`): The context tensor.
Returns:
- x (:obj:`torch.Tensor`): The output tensor after GLU transformation.
"""
gate = self.layer1(context)
gate = torch.sigmoid(gate)
x = gate * x
x = self.layer2(x)
return x
[docs]class Swish(nn.Module):
"""
Overview:
Swish activation function, which is a smooth, non-monotonic activation function. For more details, please refer
to [Searching for Activation Functions](https://arxiv.org/pdf/1710.05941.pdf).
Interfaces:
``__init__``, ``forward``.
"""
[docs] def __init__(self):
"""
Overview:
Initialize the Swish module.
"""
super(Swish, self).__init__()
[docs] def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Overview:
Compute the Swish transformation of the input tensor.
Arguments:
- x (:obj:`torch.Tensor`): The input tensor.
Returns:
- x (:obj:`torch.Tensor`): The output tensor after Swish transformation.
"""
return x * torch.sigmoid(x)
[docs]class GELU(nn.Module):
"""
Overview:
Gaussian Error Linear Units (GELU) activation function, which is widely used in NLP models like GPT, BERT.
For more details, please refer to the original paper: https://arxiv.org/pdf/1606.08415.pdf.
Interfaces:
``__init__``, ``forward``.
"""
[docs] def __init__(self):
"""
Overview:
Initialize the GELU module.
"""
super(GELU, self).__init__()
[docs] def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Overview:
Compute the GELU transformation of the input tensor.
Arguments:
- x (:obj:`torch.Tensor`): The input tensor.
Returns:
- x (:obj:`torch.Tensor`): The output tensor after GELU transformation.
"""
return 0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0))))
[docs]def build_activation(activation: str, inplace: bool = None) -> nn.Module:
"""
Overview:
Build and return the activation module according to the given type.
Arguments:
- activation (:obj:`str`): The type of activation module, now supports \
['relu', 'glu', 'prelu', 'swish', 'gelu', 'tanh', 'sigmoid', 'softplus', 'elu', 'square', 'identity'].
- inplace (Optional[:obj:`bool`): Execute the operation in-place in activation, defaults to None.
Returns:
- act_func (:obj:`nn.module`): The corresponding activation module.
"""
if inplace is not None:
assert activation == 'relu', 'inplace argument is not compatible with {}'.format(activation)
else:
inplace = False
act_func = {
'relu': nn.ReLU(inplace=inplace),
'glu': GLU,
'prelu': nn.PReLU(),
'swish': Swish(),
'gelu': GELU(),
"tanh": nn.Tanh(),
"sigmoid": nn.Sigmoid(),
"softplus": nn.Softplus(),
"elu": nn.ELU(),
"silu": torch.nn.SiLU(inplace=inplace),
"square": Lambda(lambda x: x ** 2),
"identity": Lambda(lambda x: x),
}
if activation.lower() in act_func.keys():
return act_func[activation]
else:
raise KeyError("invalid key for activation: {}".format(activation))