Class StackedLinearLayers

cdslib/core/nn/modules/linear.py:105–300

Convenient helper nn.Module to create a stack of linear layers.

class StackedLinearLayers(nn.Module):
    """
    Convenient helper nn.Module to create a stack of linear layers.
    """

    def __init__(
        self,
        num_layers: int,
        dim_input: int,
        dim_output: int,
        dim_features: T.Union[T.Sequence[int], int],
        nonlinearity: str = "leaky_relu",
        add_norm_layer: bool = False,
        norm_fun: T.Callable = nn.LayerNorm,
        dropout_prob: float = 0.0,
        output_add_nonlinearity: bool = False,
    ):
        """
        Convenient helper nn.Module to create a stack of linear layers.

        Args:
            num_layers:
                Total number of linear layers to create.
            dim_input:
                Feature dimension of the input tensor, which is :math:`(*, C_{in})`.
            dim_output:
                Feature dimension of the output tensor, which is :math:`(*, C_{out})`.
            dim_features:
                An integer if all layers share the same feature dimension,
                or a list of num_layer-1 integers, one for each layer except the last layer.
            nonlinearity:
                Nonlinearity used after each linear layer (except the last layer if output_add_nonlinearity is False).
                Choose from: `leaky_relu`, `relu`, `tanh`, `sigmoid`, `silu`, `swish`
                Note that silu (swish) is supported in pytorch version >= 1.7.0.
            add_norm_layer:
                Whether to add normalization layers between linear layers
            norm_fun:
                Callable function used to normalize the output of linear layer (before nonlinearity).
                It should be a function that takes dim_feature as input.
                For example, you can pass `nn.LayerNorm`.
                If you want to control additional functionality like the eps and elementwise_affine of nn.LayerNorm,
                you can pass a lambda function:
                lambda dim: torch.nn.LayerNorm(dim, eps=1e-5, elementwise_affine=False)
            dropout_prob:
                Dropout probability added after nonlinearity. If 0, no dropout layer is added.
            output_add_nonlinearity:
                Whether to add nonlinearity (norm_layer, and dropout) at the last layer

        Note that the order of the layers is:
            Linear -> normalization (if add_norm_layer) -> nonlinearity -> dropout.

        """
        super().__init__()
        self.dim_input = dim_input
        self.num_layers = num_layers
        self.dim_output = dim_output
        self.add_norm_layer = add_norm_layer
        self.linear_bias = not self.add_norm_layer  # if added normalization layer, no need to learn bias