| 10 | |
| 11 | |
| 12 | class MLP(nn.Module): |
| 13 | def __init__(self, |
| 14 | input_dim: int, |
| 15 | hidden_dims: Sequence[int], |
| 16 | output_dim: int, |
| 17 | net_normalization: Optional[str] = None, |
| 18 | activation_function: str = 'gelu', |
| 19 | dropout: float = 0.0, |
| 20 | residual: bool = False, |
| 21 | output_normalization: bool = False, |
| 22 | output_activation_function: Optional[Union[str, bool]] = None, |
| 23 | out_layer_bias_init: Tensor = None, |
| 24 | name: str = "" |
| 25 | ): |
| 26 | """ |
| 27 | Args: |
| 28 | input_dim (int): the expected 1D input tensor dim |
| 29 | output_activation_function (str, bool, optional): By default no output activation function is used (None). |
| 30 | If a string is passed, is must be the name of the desired output activation (e.g. 'softmax') |
| 31 | If True, the same activation function is used as defined by the arg `activation_function`. |
| 32 | """ |
| 33 | |
| 34 | super().__init__() |
| 35 | self.name = name |
| 36 | hidden_layers = [] |
| 37 | dims = [input_dim] + list(hidden_dims) |
| 38 | for i in range(1, len(dims)): |
| 39 | hidden_layers += [MLP_Block( |
| 40 | in_dim=dims[i - 1], |
| 41 | out_dim=dims[i], |
| 42 | net_norm=net_normalization.lower() if isinstance(net_normalization, str) else 'none', |
| 43 | activation_function=activation_function, |
| 44 | dropout=dropout, |
| 45 | residual=residual |
| 46 | )] |
| 47 | self.hidden_layers = nn.ModuleList(hidden_layers) |
| 48 | |
| 49 | out_weight = nn.Linear(dims[-1], output_dim, bias=True) |
| 50 | if out_layer_bias_init is not None: |
| 51 | log.info(' Pre-initializing the MLP final/output layer bias.') |
| 52 | out_weight.bias.data = out_layer_bias_init |
| 53 | out_layer = [out_weight] |
| 54 | if output_normalization and net_normalization != 'none': |
| 55 | out_layer += [get_normalization_layer(net_normalization, output_dim)] |
| 56 | if output_activation_function is not None and output_activation_function: |
| 57 | if isinstance(output_activation_function, bool): |
| 58 | output_activation_function = activation_function |
| 59 | |
| 60 | out_layer += [get_activation_function(output_activation_function, functional=False)] |
| 61 | self.out_layer = nn.Sequential(*out_layer) |
| 62 | |
| 63 | def forward(self, X: Tensor) -> Tensor: |
| 64 | for layer in self.hidden_layers: |
| 65 | X = layer(X) |
| 66 | |
| 67 | Y = self.out_layer(X) |
| 68 | return Y.squeeze(1) |
| 69 | |