Initialize weights in PyTorch

If we are trying to build a neural network then we have to initialize the layers of the network with some initial weights which we try to optimize as the training process of the model goes on. The method by which the weights of a neural network are initialized does affect the time required to reach the optimized solution and solve the problem of vanishing or exploding gradients. In this article, we will try to learn the method by which effective initialization of weights can be done by using the PyTorch machine learning framework.

Why initialize weights?

Initializing the weights of a neural network is a vital step in the training process as appropriate weight initialization is an instrumental factor impacting the convergence and performance of a network. Weights that are initialized to the same value can cause the model to converge to the same suboptimal solution, regardless of the optimization algorithm being used.

Weights that are initialized to large values can lead to vanishing or exploding gradients, depending on the activation function being used. This can cause the model to converge slowly or not at all. Weights that are initialized to small random values can lead to more efficient training, as the optimization algorithm is able to make larger updates to the weights at the beginning of training. Different initialization methods can be more suitable for different types of problems and model architectures.

Using the nn.init Module for Weights Initialization

The PyTorch nn.init module is a conventional way to initialize weights in a neural network, which provides a multitude of weight initialization methods such as:

1. Uniform initialization
2. Xavier initialization
3. Kaiming initialization
4. Zeros initialization
5. One’s initialization
6. Normal initialization

An example implementation of the same is provided below:

Uniform Initialization

Using a uniform distribution to initialize the weights can help prevent the ‘vanishing gradient’ problem, as the distribution has a finite range and the weights are distributed evenly across that range. However, this method can suffer from the ‘exploding gradient’ problem if the range is too large.

Python3

import torch

# Initializing a linear layer with 2
# independent features and 3 dependent features
linear_layer = torch.nn.Linear(2, 3)

# Initializing the weights with a uniform distribution
torch.nn.init.uniform_(linear_layer.weight)

# Displaying the initialized weights
print(linear_layer.weight)

Output:

Parameter containing:
tensor([[-0.1768, -0.4942],
       [ 0.0756, -0.0967],
       [-0.3923,  0.3283]], requires_grad=True)

Xavier Initialization

Using Xavier initialization can help prevent the ‘vanishing gradient’ problem, as it scales the weights such that the variance of the outputs of each layer is the same as the variance of the inputs.

Python3

import torch

# Initializing a linear layer with
# 2 independent features and 3 dependent features
linear_layer = torch.nn.Linear(2, 3)

# Initializing the weights with the Xavier initialization method
torch.nn.init.xavier_uniform_(linear_layer.weight)

# Displaying the initialized weights
print(linear_layer.weight)

Output:

Parameter containing:
tensor([[ 0.4442, -0.3890],
        [-0.2876, -0.3379],
        [-0.5261,  0.5227]], requires_grad=True)

Kaiming Initialization

Using Kaiming initialization can help prevent the ‘vanishing gradient’ problem, as it scales the weights such that the variance of the outputs is the same as the variance of the inputs, taking into account the nonlinearity of the activation function.

Python3

import torch

# Initializing a linear layer with
# 2 independent features and 3 dependent features
linear_layer = torch.nn.Linear(2, 3)

# Initializing the weights with the Kaiming initialization method
torch.nn.init.kaiming_uniform_(linear_layer.weight,
                               a=0, mode="fan_in",
                               nonlinearity="relu")

# Displaying the initialized weights
print(linear_layer.weight)

Output:

Parameter containing:
tensor([[ 0.0582,  0.4701],
        [ 0.4982,  0.5452],
        [-0.0384,  0.5999]], requires_grad=True)

Zeros and Ones Initialisation

Initializing the weights to zeros can cause the model to converge slowly, as all of the weights will be updated in the same direction. This can also lead to the ‘vanishing gradient’ problem. 

Python3

import torch

# Initializing a linear layer with
# 2 independent features and 3 dependent features
linear_layer = torch.nn.Linear(2, 3)

# Initializing the weights with the
# zeros initialization method
torch.nn.init.zeros_(linear_layer.weight)

# Displaying the initialized weights
print(linear_layer.weight)

Output:

Parameter containing:
tensor([[0., 0.],
        [0., 0.],
        [0., 0.]], requires_grad=True)

Initializing the weights to ones can cause the model to converge slowly, as all of the weights will be updated in the same direction. This can also lead to the ‘exploding gradient’ problem. 

Python3

import torch

# Initializing a linear layer with
# 2 independent features and 3 dependent features
linear_layer = torch.nn.Linear(2, 3)

# Initializing the weights with the
# ones initialization method
torch.nn.init.ones_(linear_layer.weight)

# Displaying the initialized weights
print(linear_layer.weight)

Output:

Parameter containing:
tensor([[1., 1.],
        [1., 1.],
        [1., 1.]], requires_grad=True)

Normal Initialisation

Using a normal distribution to initialize the weights can help prevent the ‘exploding gradient’ problem, as the distribution has a finite range and the weights are distributed evenly around the mean. It must be noted that the neural network’s performance is not impacted by the weights alone; the learning rate, the optimization algorithms and the hyperparameters used also play a crucial role in increasing the efficiency of the neural network.

Python3

import torch

# Initializing a linear layer with
# 2 independent features and 3 dependent features
linear_layer = torch.nn.Linear(2, 3)

# Initializing the weights with the
# normal initialization method
torch.nn.init.normal_(linear_layer.weight,
                      mean=0, std=1)

# Displaying the initialized weights
print(linear_layer.weight)

Output:

Parameter containing:
tensor([[-0.1759,  0.5192],
        [-0.5621, -0.3871],
        [-0.6071,  0.3538]], requires_grad=True)

Applying a Custom Function for Weights Initialization

An alternative method is to create a customized function to initialize the weights, which can be applied to the layer using the apply attribute.

Python3

import torch

# User defined function to initialize the weights
def custom_weights(m):
    torch.nn.init.uniform_(m.weight,
                           -0.5, 0.5)

# Initializing a linear layer with
# 2 independent features and 3 dependent features
linear_layer = torch.nn.Linear(2, 3)

# Applying the user defined function to the layer
linear_layer.apply(custom_weights)

# Displaying the initialized weights
print(linear_layer.weight)

Output:

Parameter containing:
tensor([[ 0.4341, -0.3424],
        [ 0.2095,  0.1782],
        [-0.4244,  0.1719]], requires_grad=True)
        

Using a user-defined Layer Class for Weights Initialization

Another method involves creating a user-defined class that inherits from the torch.nn.Module class. Therein, the constructor can be overridden in order to implement custom weights.

Python3

import torch

# User defined Layer
class MyLayer(torch.nn.Module):

    # Overriding the constructor
    def __init__(self, independent, dependent):
        # Calling the super-class' constructor
        super(MyLayer, self).__init__()
        self.linear = torch.nn.Linear(independent,
                                      dependent)
        torch.nn.init.uniform_(self.linear.weight,
                               -0.5, 0.5)

    def forward(self, x):
        return self.linear(x)


# Initializing a linear layer with
# 2 independent features and 3 dependent features
linear_layer = MyLayer(2, 3)

# Displaying the initialized weights
print(linear_layer.linear.weight)

Output:

Parameter containing:
tensor([[-0.1566,  0.2461],
        [-0.3361, -0.0551],
        [ 0.4607,  0.3077]], requires_grad=True)

In conclusion, initializing the weights of a neural network model is an important step in the training process, as it can have a significant impact on the model’s performance. PyTorch provides several built-in initialization methods, including uniform, normal, Xavier, Kaiming, ones, and zeros. Each of these methods has its own advantages and disadvantages, and the choice of method will depend on the specific problem and model architecture being used. It is important to choose an initialization method that is suitable for the problem at hand, as it can help prevent vanishing or exploding gradient problems and improve the convergence speed and final accuracy of the model.