Weight initialization techniques such as Xavier or He initialization are used to set the initial weights of a neural network in order to improve its performance and convergence during training. In PyTorch, these techniques can be implemented easily using the ‘torch.nn.init‘ module.
Xavier initialization, also known as Glorot initialization, sets the initial weights to random values drawn from a uniform distribution that is dependent on the number of input and output neurons in a layer. This can be achieved using the ‘torch.nn.init.xavier_uniform_()‘ or ‘torch.nn.init.xavier_normal_()‘ functions, depending on whether a uniform or normal distribution is desired. For example, to initialize the weights of a linear layer using Xavier uniform distribution, you can use the following code:
import torch
import torch.nn as nn
layer = nn.Linear(in_features=10, out_features=20)
torch.nn.init.xavier_uniform_(layer.weight)He initialization, on the other hand, sets the initial weights to random values drawn from a normal distribution with zero mean and a variance that is dependent on the number of input neurons in a layer. This can be achieved using the ‘torch.nn.init.kaiming_uniform_()‘ or ‘torch.nn.init.kaiming_normal_()‘ functions, depending on whether a uniform or normal distribution is desired. For example, to initialize the weights of a convolutional layer using He normal distribution, you can use the following code:
import torch
import torch.nn as nn
layer = nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3)
torch.nn.init.kaiming_normal_(layer.weight)It is worth noting that these initialization techniques are typically applied to the weights of hidden layers, as the initial weights of output layers are often initialized using a procedure specific to the task being tackled (e.g. a softmax function for classification tasks).