Understand Kaiming Initialization and Implementation Detail in PyTorch | by Xu LIANG | Towards Data Science

If you create weight implicitly by creating a linear layer, you should set modle='fan_in'.
linear = torch.nn.Linear(node_in, node_out)init.kaiming_normal_(linear.weight, mode=’fan_in’)t = relu(linear(x_valid))
If you create weight explicitly by creating a random matrix, you should set modle='fan_out'.
w1 = torch.randn(node_in, node_out)init.kaiming_normal_(w1, mode=’fan_out’)b1 = torch.randn(node_out)t = relu(linear(x_valid, w1, b1))
The content is structured as follows.
Weight Initialization Matters!What is Kaiming initialization?Why Kaiming initialization works?Understand fan_in and fan_out mode in Pytorch implementation
Weight Initialization Matters!
What is Kaiming initialization?
Why Kaiming initialization works?
Understand fan_in and fan_out mode in Pytorch implementation
Initialization is a process to create weight. In the below code snippet, we create a weight w1 randomly with the size of(784, 50).
torhc.randn(*sizes) returns a tensor filled with random numbers from a normal distribution with mean 0 and variance 1 (also called the standard normal distribution). The shape of the tensor is defined by the variable argument sizes.
And this weight will be updated during the training phase.
# random initw1 = torch.randn(784, 50) b1 = torch.randn(50)def linear(x, w, b):    return x@w + bt1 = linear(x_valid, w1, b1)print(t1.mean(), t1.std())############# output ##############tensor(3.5744) tensor(28.4110)
You may wonder why need we care about initialization if the weight can be updated during the training phase. No matter how to initialize the weight, it will be updated “well” eventually.
But the reality is not so sweet. If we random initialize the weight, it will cause two problems, the vanishing gradient problem and exploding gradient problem.
Vanishing gradient problem means weights vanish to 0. Because these weights are multiplied along with the layers in the backpropagation phase. If we initialize weights very small(<1), the gradients tend to get smaller and smaller as we go backward with hidden layers during backpropagation. Neurons in the earlier layers learn much more slowly than neurons in later layers. This causes minor weight updates.
Exploding gradient problem means weights explode to infinity(NaN). Because these weights are multiplied along with the layers in the backpropagation phase. If we initialize weights very large(>1), the gradients tend to get larger and larger as we go backward with hidden layers during backpropagation. Neurons in the earlier layers update in huge steps, W = W — ⍺ * dW, and the downward moment will increase.
Kaiming et al. derived a sound initialization method by cautiously modeling non-linearity of ReLUs, which makes extremely deep models (>30 layers) to converge. Below is the Kaiming initialization function.
a: the negative slope of the rectifier used after this layer (0 for ReLU by default)
fan_in: the number of input dimension. If we create a (784, 50), the fan_in is 784. fan_in is used in the feedforward phase. If we set it as fan_out, the fan_out is 50. fan_out is used in the backpropagation phase. I will explain two modes in detail later.
We compare the random initialization and Kaiming initialization to show the effectiveness of Kaiming initialization.
Random Initialization
# random initw1 = torch.randn(784, 50) b1 = torch.randn(50)w2 = torch.randn(50, 10) b2 = torch.randn(10)w3 = torch.randn(10, 1) b3 = torch.randn(1)def linear(x, w, b):    return x@w + bdef relu(x):    return x.clamp_min(0.)t1 = relu(linear(x_valid, w1, b1))t2 = relu(linear(t1, w2, b2))t3 = relu(linear(t2, w3, b3))print(t1.mean(), t1.std())print(t2.mean(), t2.std())print(t3.mean(), t3.std())############# output ##############tensor(13.0542) tensor(17.9457)tensor(93.5488) tensor(113.1659)tensor(336.6660) tensor(208.7496)
We initialize weight with a normal distribution with mean 0 and variance 1, and the ideal distribution of weight after ReLU should have slightly incremented mean layer by layer and variance close to 1. But the distribution changes a lot after some layers in the feedforward phase.
Why the mean of weight should be slightly incremented layer by layer?
Because we use the ReLU as the activation function. ReLU will return the value provided if input value is bigger than 0 and return value 0 if the input value is less than 0.
if input < 0:    return 0else:    return input
After ReLU, all negative values become 0. The mean will become larger when the layer becomes deeper.
Kaiming Initialization
# kaiming initnode_in = 784node_out = 50# random initw1 = torch.randn(784, 50) * math.sqrt(2/784)b1 = torch.randn(50)w2 = torch.randn(50, 10) * math.sqrt(2/50)b2 = torch.randn(10)w3 = torch.randn(10, 1) * math.sqrt(2/10)b3 = torch.randn(1)def linear(x, w, b):    return x@w + bdef relu(x):    return x.clamp_min(0.)t1 = relu(linear(x_valid, w1, b1))t2 = relu(linear(t1, w2, b2))t3 = relu(linear(t2, w3, b3))print(t1.mean(), t1.std())print(t2.mean(), t2.std())print(t3.mean(), t3.std())############# output ##############tensor(0.7418) tensor(1.0053)tensor(1.3356) tensor(1.4079)tensor(3.2972) tensor(1.1409)
We initialize weight with a normal distribution with mean 0 and variance std, and the ideal distribution of weight after relu should have slightly incremented mean layer by layer and variance close to 1. We can see the output is close to what we expected. The mean increment slowly and std is close to 1 in the feedforward phase. And such stability will avoid the vanishing gradient problem and exploding gradient problem in the backpropagation phase.
Kaiming initialization shows better stability than random initialization.
nn.init.kaiming_normal_() will return tensor that has values sampled from mean 0 and variance std. There are two ways to do it.
One way is to create weight implicitly by creating a linear layer. We set mode='fan_in' to indicate that using node_in calculate the std
from torch.nn import init# linear layer implementationnode_in, node_out = 784, 50layer = torch.nn.Linear(node_in, node_out)init.kaiming_normal_(layer.weight, mode='fan_in')t = relu(layer(x_valid))print(t.mean(), t.std())############# output ##############tensor(0.4974, grad_fn=<MeanBackward0>) tensor(0.8027, grad_fn=<StdBackward0>)
Another way is to create weight explicitly by creating a random matrix, you should set mode='fan_out'.
def linear(x, w, b):    return x@w + b# weight matrix implementationnode_in, node_out = 784, 50w1 = torch.randn(node_in, node_out)init.kaiming_normal_(w1, mode='fan_out')b1 = torch.randn(node_out)t = relu(linear(x_valid, w1, b1))print(t.mean(), t.std())############# output ##############tensor(0.6424) tensor(0.9772)
Two implementation methods are both right. The mean is close to 0.5 and std is close to 1. But wait a minute, do you find something strange?
Why the mode is different?
According to the document, choosing 'fan_in' preserves the magnitude of the variance of the weights in the forward pass. Choosing 'fan_out' preserves the magnitudes in the backward pass. We can write as below.
node_in, node_out = 784, 50# fan_in modeW = torch.randn(node_in, node_out) * math.sqrt(2 / node_in)# fan_out modeW = np.random.randn(node_in, node_out) * math.sqrt(2/ node_out)
In the linear layer implementation, we set mode='fan_in'. Yes, this is the feedforward phase, we should set mode='fan_in' . Nothing wrong.
But why we set the mode as fan_out in the weight matrix implementation?
The reason behind the source code of nn.init.kaiming_normal_()
def _calculate_fan_in_and_fan_out(tensor):    dimensions = tensor.dim()    if dimensions < 2:        raise ValueError("Fan in and fan out can not be computed for tensor with fewer than 2 dimensions")if dimensions == 2:  # Linear        fan_in = tensor.size(1)        fan_out = tensor.size(0)    else:        num_input_fmaps = tensor.size(1)        num_output_fmaps = tensor.size(0)        receptive_field_size = 1        if tensor.dim() > 2:            receptive_field_size = tensor[0][0].numel()        fan_in = num_input_fmaps * receptive_field_size        fan_out = num_output_fmaps * receptive_field_sizereturn fan_in, fan_out
This is the source code to get the right mode. The tensor is the w1 with size of (784, 50). So fan_in = 50, fan_out=784. When we set the mode as fan_out in the weight matrix implementation. The init.kaiming_normal_() actually calculates like below.
node_in, node_out = 784, 50W = np.random.randn(node_in, node_out)init.kaiming_normal_(W, mode='fan_out')# what init.kaiming_normal_() actually does    # fan_in = 50    # fan_out = 784W = W * torch.sqrt(784 / 2)
Ok, make sense. But how to explain using fan_in in the linear layer implementation?
When we use linear to create weight implicitly, the weight is transposed implicitly. Here is the source code of torch.nn.functional.linear.
def linear(input, weight, bias=None):    # type: (Tensor, Tensor, Optional[Tensor]) -> Tensor    r"""    Applies a linear transformation to the incoming data: :math:`y = xA^T + b`.    Shape:        - Input: :math:`(N, *, in\_features)` where `*` means any number of          additional dimensions        - Weight: :math:`(out\_features, in\_features)`        - Bias: :math:`(out\_features)`        - Output: :math:`(N, *, out\_features)`    """    if input.dim() == 2 and bias is not None:        # fused op is marginally faster        ret = torch.addmm(bias, input, weight.t())    else:        output = input.matmul(weight.t())        if bias is not None:            output += bias        ret = output    return ret
The weight is initialized with the size of (out_features, in_features). For example, if we input the size (784, 50) , the size of weight is actually (50, 784).
torch.nn.Linear(784, 50).weight.shape############# output ##############torch.Size([50, 784])
That’s why linear need to first transpose the weight and then do the matmul operation.
    if input.dim() == 2 and bias is not None:        # fused op is marginally faster        ret = torch.addmm(bias, input, weight.t())    else:        output = input.matmul(weight.t())
Because the weight in the linear layer has size of (50, 784), the init.kaiming_normal_() actually calculates like below.
node_in, node_out = 784, 50layer = torch.nn.Linear(node_in, node_out)init.kaiming_normal_(layer.weight, mode='fan_out')# the size of layer.weight is (50, 784)# what init.kaiming_normal_() actually does    # fan_in = 784    # fan_out = 50W = W * torch.sqrt(784 / 2)
In this post, I first talked about why initialization matters and what is kaiming initialization. And I break down how to use PyTorch to implement it. Hope this post be helpful. Leave a comment if you have any advice.
The full code in this snippet.
Check out my other posts on Medium with a categorized view!GitHub: BrambleXuLinkedIn: Xu LiangBlog: BrambleXu
Why cautiously initializing deep neural networks matters?
Deep Learning Best Practices (1) — Weight Initialization
Kaiming Initialization paper: Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification
A Gentle Introduction to the Rectified Linear Unit (ReLU)
Fast.ai’s course Deep Learning for coders lesson 8