Class StackedModulatedSubspace

cdslib/core/nn/modules/subspace.py:315–531

Convenient helper nn.Module to create a stack of modulated subspace layers.

class StackedModulatedSubspace(nn.Module):
    """
    Convenient helper nn.Module to create a stack of modulated subspace layers.
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        layer_configs: T.Sequence[T.Dict[str, T.Any]],
    ):
        r"""Construct multiple stacked modulated subspace layers
        sandwiched between nonlinearity.

        Args:
            in_features (int):
                number of input channels
            out_features (int):
                number of output channels in the output system
            layer_configs (list of T.Dict[str, T.Any]):
                a list of dict (one for each layer) containing the parameters
                to :py:class:`ModulatedSubspace`
                and parameters for nonlinearity, and dropout:

                - out_features (int, required):
                    output feature dimension
                - style_features (int, required):
                    style feature dimension
                - bias (bool):
                    whether to learn bias :math:`b`. Default: `True`.
                - fixed_bias (float):
                    a fixed bias b0 added after :math:`Wx + b + b0`. Default: `0`.
                - input_bias (bool):
                    whether to learn x0. Default: `True`.
                - fixed_input_bias (float):
                    a fixed bias x1 added after x0. Default: `0`.
                - lr_multiplier (float):
                    a factor controls the learning rate of the layer. Default: `1`.
                - demodulate (bool):
                    whether to normalize the row of W. Default: `True`.
                - normalize_basis (bool):
                    whether to normalize the basis to have unit l2 norm. Default: `True`.
                - orthogonalize_basis (bool):
                    whether to orthogonalize the basis. Default: `False`.
                    (recommended: False)
                - dropout (float):
                    the dropout rate. Default: `0`.
                - nonlinearity (str):
                    ``'none'``, ``'leaky_relu'``, ``'relu'``, ``'tanh'``,
                    ``'sigmoid'``, ``'silu'``. Default: ``'relu'``.

        """
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.layer_configs = layer_configs
        self.num_layers = len(self.layer_configs)
        assert self.out_features == self.layer_configs[-1]["out_features"]