Class LinearLayer

cdslib/core/nn/modules/linear.py:18–102

A helper class that wraps around the typical nn.Linear to help initialize it.

class LinearLayer(nn.Linear):
    """
    A helper class that wraps around the typical nn.Linear to help initialize it.
    """

    _FLOAT_MODULE = nn.Linear

    def __init__(
        self,
        in_features: int,
        out_features: int,
        bias: bool = True,
        w_init_gain: str = "linear",
        init_method: str = "xavier_normal",
        lrelu_nslope: float = 0.01,
        kaiming_fan_mode: str = "fan_in",
        bias_init_val: float = 0.0,
        lr_multiplier: float = 1.0,
        fixed_bias: float = None,
    ):
        """
        Create a linear layer whose weights are initialized with a chosen method.
        The usage of the layer is the same as torch.nn.Linear.

        Args:
            in_features:
                Size of the last dimension of the input tensor.
            out_features:
                Size of the last dimension of the output tensor.
            bias:
                Whether to add a learnable bias term, one for each out_features.
            w_init_gain:
                The nonlinearity that is designed to be added after the layer.
                Check :py:func:`cdslib.nn.init_weight` for supported nonlinearities.
                Note that the layer does not add nonlinearity.
            init_method:
                The initialization method. Check :py:func:`cdslib.nn.init_weight`.
            lrelu_nslope:
                The negative slope of leaky-relu if leaky-relu is used. Check :py:func:`cdslib.nn.init_weight`.
            kaiming_fan_mode:
                The fan mode if kaiming_* init method is used. Check :py:func:`cdslib.nn.init_weight`.
            bias_init_val:
                Initialization value of the bias.
            lr_multiplier:
                A scalar to be multiplied with the learning rate.
                For example, if lr_multiplier = 0.1, both the weight and the bias will update at 1/10 of the rate.
            fixed_bias (float):
                A fixed bias value to add to the output. If `None`, nothing is added.
        """
        nn.Linear.__init__(self, in_features, out_features, bias)
        nn_utils.init_weight(
            self.weight,
            w_init_gain=w_init_gain,  # leaky_relu, relu, tanh, sigmoid, ...
            init_method=init_method,
            lrelu_nslope=lrelu_nslope,
            kaiming_fan_mode=kaiming_fan_mode,
        )
        # init bias