Function init_weight

cdslib/core/nn/nn_utils.py:17–100

A helper function to initialize the weights of a linear/convolutional layer. Args: weight: (*), an n-dimensional torch.Tensor to be initialized w_init_gain: The nonlinearity after the linear layer. Can be chosen from the functions supported by

(
    weight: torch.Tensor,
    w_init_gain: str = "linear",
    init_method: str = "xavier_normal",
    lrelu_nslope: float = 0.01,
    kaiming_fan_mode: str = "fan_in",
)

def init_weight(
    weight: torch.Tensor,
    w_init_gain: str = "linear",
    init_method: str = "xavier_normal",
    lrelu_nslope: float = 0.01,
    kaiming_fan_mode: str = "fan_in",
):
    """
    A helper function to initialize the weights of a linear/convolutional layer.

    Args:
        weight:
            (*), an n-dimensional torch.Tensor to be initialized
        w_init_gain:
            The nonlinearity after the linear layer. Can be chosen from the functions supported by
            torch.nn.init.calculate_gain.
            This includes:
            'linear', 'relu', 'silu', 'leaky_relu', 'relu', 'tanh', 'sigmoid'.
        init_method:
            The initialization method. Can be chosen from:
                'normal':
                    randomly sampeld from a Gaussian distribution
                'uniform':
                    randomly sampeld from a uniform distribution
                'xavier_uniform':
                    Check :py:func:`torch.nn.init.xavier_uniform_`.
                'xavier_normal'
                    Check :py:func:`torch.nn.init.xavier_normal_`.
                'xavier':
                    same as 'xavier_normal'
                'kaiming_uniform':
                    Check :py:func:`torch.nn.init.kaiming_uniform_`.
                'kaiming_normal'
                    Check :py:func:`torch.nn.init.kaiming_normal_`.
                'kaiming':
                    same as 'kaiming_normal'
                'orthogonal':
                    Check :py:func:`torch.nn.init.orthogonal_`.
        lrelu_nslope:
            Negative slope used in the leaky-relu.
        kaiming_fan_mode:
            Fan mode used by kaiming_* init_methods.
    Returns:
        Does not return. The function directly modifies the content of weight.

    Note that the function contains torch.no_grad, so there is no need to wrap it with one.
    """

    # handle silu/swish
    if w_init_gain in {"silu", "swish"}:
        # since silu and relu has similar shape, use the gain for relu
        w_init_gain = "relu"

    # calculate gain
    if w_init_gain == "leaky_relu":
        gain = torch.nn.init.calculate_gain(w_init_gain, lrelu_nslope)
        kaiming_a = lrelu_nslope
    else: