mlp

mlp

MLP

MLP(input_dimension, hidden_dimensions=None, output_dim=None, activation_function=GELU, dropout=0.0)

Bases: Module

Multi-layer perceptron with configurable hidden sizes, activation, and dropout.

Multi-layer Perceptron (MLP) module.

Parameters:

Name	Type	Description	Default
input_dimension	int	Input feature dimension	required
hidden_dimensions	list[int] | None	List of hidden layer dimensions	None
output_dim	int | None	Output feature dimension	None
activation_function	Callable	Activation function class callable	GELU
dropout	float	Dropout rate between layers	0.0

Source code in src/versatil/models/layers/mlp.py

def __init__(
    self,
    input_dimension: int,
    hidden_dimensions: list[int] | None = None,
    output_dim: int | None = None,
    activation_function: Callable = nn.GELU,
    dropout: float = 0.0,
):
    """Multi-layer Perceptron (MLP) module.

    Args:
        input_dimension: Input feature dimension
        hidden_dimensions: List of hidden layer dimensions
        output_dim: Output feature dimension
        activation_function: Activation function class callable
        dropout: Dropout rate between layers
    """
    super().__init__()
    hidden_dimensions = hidden_dimensions if hidden_dimensions is not None else []
    layers: list[nn.Module] = []
    prev_dim = input_dimension
    for hidden_dimension in hidden_dimensions:
        if issubclass(activation_function, GatedLinearUnit):
            layers.append(
                activation_function(
                    input_dimension=prev_dim, hidden_dimension=hidden_dimension
                )
            )
        else:
            layers.append(nn.Linear(prev_dim, hidden_dimension))
            layers.append(activation_function())
        if dropout > 0.0:
            layers.append(nn.Dropout(dropout))
        prev_dim = hidden_dimension
    if output_dim is not None:
        layers.append(nn.Linear(prev_dim, output_dim))
    self.layers = nn.Sequential(*layers)

forward

forward(x)

Forward pass through the MLP.

Source code in src/versatil/models/layers/mlp.py

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Forward pass through the MLP."""
    result: torch.Tensor = self.layers(x)
    return result