neuraltrain.models.common.Mlp

class neuraltrain.models.common.Mlp(*, input_size: int | None = None, hidden_sizes: list[int] | None = None, norm_layer: Literal['layer', 'batch', 'instance', 'unit', None] = None, activation_layer: Literal['relu', 'gelu', 'elu', 'prelu', None] = 'relu', bias: bool = True, dropout: float = 0.0)

Multilayer perceptron, e.g. for use as projection head.

Notes

Input size can be specified in the config or at build time. Output size can be specified at build time through the output_size parameter: it will be appended to hidden_sizes as the final layer. When hidden_sizes is empty or None and output_size is given, build() returns a single nn.Linear layer. When both are absent it returns nn.Identity.