mixmo.utils.torchutils.kaiming_normal_truncated

mixmo.utils.torchutils.kaiming_normal_truncated(tensor, a=0, mode='fan_in', nonlinearity='relu')[source]

Fills the input Tensor with values according to the method described in Delving deep into rectifiers: Surpassing human-level performance on ImageNet classification - He, K. et al. (2015), using a normal distribution. The resulting tensor will have values sampled from \(\mathcal{N}(0, \text{std})\) where

\[\text{std} = \sqrt{\frac{2}{(1 + a^2) \times \text{fan\_in}}}\]

Also known as He initialization.

Parameters

  • tensor – an n-dimensional torch.Tensor

  • a – the negative slope of the rectifier used after this layer (0 for ReLU by default)

  • mode – either 'fan_in' (default) or 'fan_out'. Choosing 'fan_in' preserves the magnitude of the variance of the weights in the forward pass. Choosing 'fan_out' preserves the magnitudes in the backwards pass.

  • nonlinearity – the non-linear function (nn.functional name), recommended to use only with 'relu' or 'leaky_relu' (default).

Examples

>>> w = torch.empty(3, 5)
>>> nn.init.kaiming_normal_(w, mode='fan_out', nonlinearity='relu')