Spaces:
Running
Running
File size: 1,870 Bytes
5769ee4 |
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 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 |
from collections import OrderedDict
import torch
import torch.nn as nn
class MLP(nn.Module):
"""Basic MLP implementation with FC layers and ReLU activation.
Args:
input_dim : dimension of the input variable
output_dim : dimension of the output variable
h_dim : dimension of a hidden layer of MLP
num_h_layers : number of hidden layers in MLP
add_residual : set to True to add input to output (res-net) set to False to have pure MLP
"""
def __init__(
self,
input_dim: int,
output_dim: int,
h_dim: int,
num_h_layers: int,
add_residual: bool,
) -> None:
super().__init__()
self._input_dim = input_dim
self._output_dim = output_dim
self._h_dim = h_dim
self._num_h_layers = num_h_layers
layers = OrderedDict()
if num_h_layers > 0:
layers["fc_0"] = nn.Linear(input_dim, h_dim)
layers["relu_0"] = nn.ReLU()
else:
h_dim = input_dim
for ii in range(1, num_h_layers):
layers["fc_{}".format(ii)] = nn.Linear(h_dim, h_dim)
layers["relu_{}".format(ii)] = nn.ReLU()
layers["fc_{}".format(num_h_layers)] = nn.Linear(h_dim, self._output_dim)
self.mlp = nn.Sequential(layers)
if add_residual:
self.residual_layer = nn.Linear(input_dim, output_dim)
else:
self.residual_layer = lambda x: 0
self._layer_norm = nn.LayerNorm(output_dim)
def forward(self, input: torch.Tensor) -> torch.Tensor:
"""Forward function for MLP
Args:
input (torch.Tensor): (batch_size, input_dim) tensor
Returns:
torch.Tensor: (batch_size, output_dim) tensor
"""
return self._layer_norm(self.mlp(input) + self.residual_layer(input))
|