Skip to content

ocl.neural_networks

build_two_layer_mlp

Build a two layer MLP, with optional initial layer norm.

Separate class as this type of construction is used very often for slot attention and transformers.

Source code in ocl/neural_networks/convenience.py 
72
73
74
75
76
77
78
79
80
81
82
def build_two_layer_mlp(
    input_dim, output_dim, hidden_dim, initial_layer_norm: bool = False, residual: bool = False
):
    """Build a two layer MLP, with optional initial layer norm.

    Separate class as this type of construction is used very often for slot attention and
    transformers.
    """
    return build_mlp(
        input_dim, output_dim, [hidden_dim], initial_layer_norm=initial_layer_norm, residual=residual
    )