Weight Initialization in PyTorch

Posted by on

Weight initialization is an important technique in deep learning that involves setting the initial values of the weights in a neural network. The goal of weight initialization is to set the initial weights in such a way that the network converges faster and more accurately during training.

In PyTorch, weight initialization can be done using the torch.nn.init module, which provides various functions for initializing weights in different ways. Some commonly used initialization methods are:
Uniform Initialization: This initializes the weights with a random value between a specified range. The torch.nn.init.uniform_ function can be used for this.

Normal Initialization: This initializes the weights with a random value sampled from a normal distribution with mean 0 and standard deviation 1. The torch.nn.init.normal_ function can be used for this.

Xavier Initialization: This initializes the weights in such a way that the variance of the outputs of each layer is approximately equal to the variance of its inputs. The torch.nn.init.xavier_uniform_ and torch.nn.init.xavier_normal_ functions can be used for this.

He Initialization: This initializes the weights in such a way that the variance of the outputs of each layer is equal to the variance of its inputs divided by the number of inputs. The torch.nn.init.kaiming_uniform_ and torch.nn.init.kaiming_normal_ functions can be used for this.

These initialization methods can be applied to the weights of a layer using the apply method of the layer object.

For example, to apply Xavier initialization to the weights of a linear layer, we can do the following:

import torch.nn as nn 
layer = nn.Linear(in_features=100, out_features=50) 
nn.init.xavier_uniform_(layer.weight)

Output

Parameter containing:
tensor([[ 0.1850,  0.1414, -0.0704,  ..., -0.0037,  0.1601, -0.1542],
        [-0.1447, -0.1670, -0.0843,  ..., -0.0606, -0.1963, -0.1505],
        [-0.0898,  0.1764,  0.0887,  ...,  0.1731,  0.1025,  0.1911],
        ...,
        [-0.0995, -0.1905,  0.0787,  ...,  0.0343, -0.0907, -0.1569],
        [-0.0219, -0.1083, -0.0351,  ...,  0.0601, -0.0468, -0.0675],
        [-0.1896,  0.0532, -0.0470,  ..., -0.1301, -0.0343, -0.0253]],
       requires_grad=True)

This initializes the weights of the layer object with Xavier initialization.

Example

Here's an example code for weight initialization in PyTorch:
import torch 
import torch.nn as nn 
# Define a simple neural network with 2 layers 
class Net(nn.Module): 
    def __init__(self): 
        super(Net, self).__init__() 
        self.fc1 = nn.Linear(784, 128) 
        self.fc2 = nn.Linear(128, 10) 
    def forward(self, x): 
        x = torch.flatten(x, 1) 
        x = self.fc1(x) 
        x = torch.relu(x) 
        x = self.fc2(x) 
        return x 
# Create an instance of the network 
net = Net() 
# Initialize the weights of the first layer using Xavier initialization 
nn.init.xavier_uniform_(net.fc1.weight) 
# Initialize the biases of both layers with zeros 
nn.init.zeros_(net.fc1.bias) 
nn.init.zeros_(net.fc2.bias) 
# Generate some random input data 
x = torch.randn(1, 1, 28, 28) 
# Compute the output of the network 
output = net(x) 
# Print the output tensor 
print(output)

Output

tensor([[ 0.3290, -0.2978,  0.5987, -0.0243, -1.1859, -0.4180, -0.1961, -0.2199,
         -0.4154,  0.3470]], grad_fn=<AddmmBackward0>)

This code defines a simple neural network with 2 fully connected layers and applies weight initialization to the first layer using Xavier initialization. The biases of both layers are initialized to zero using the nn.init.zeros_ function.

The code then generates some random input data and computes the output of the network. Finally, it prints the output tensor. Note that this is just a simple example and you can modify it according to your needs.

More About Xavier Initialization

Xavier initialization is a technique for initializing the weights of a neural network in a way that helps improve the performance of the network during training. It was proposed by Xavier Glorot and Yoshua Bengio in their 2010 paper "Understanding the difficulty of training deep feedforward neural networks".

The idea behind Xavier initialization is to set the initial weights of each layer in a way that the variance of the outputs of the layer is approximately equal to the variance of its inputs. This helps ensure that the activations of the neurons in the layer do not become too large or too small, which can cause the gradients to become too small or too large during backpropagation.

Xavier initialization is based on the following assumption:

Suppose we have a layer with n_in input neurons and n_out output neurons.

  • The input to each neuron in the layer is a linear combination of the inputs to the previous layer.
  • The weights connecting the previous layer to the current layer are initialized with random values drawn from a normal distribution with mean 0 and standard deviation sigma.
  • The bias terms are initialized with zeros.
Under these assumptions, Xavier initialization sets sigma to:

sigma = sqrt(2 / (n_in + n_out))
This ensures that the variance of the outputs of each neuron in the layer is approximately equal to the variance of its inputs. Note that the factor of 2 in the denominator is used to account for the fact that the activation function used in the layer is typically a "centered" function (i.e., it has a mean of 0).

In PyTorch, you can use the torch.nn.init.xavier_uniform_ or torch.nn.init.xavier_normal_ functions to apply Xavier initialization to the weights of a layer. The xavier_uniform_ function initializes the weights with values drawn from a uniform distribution, while the xavier_normal_ function initializes the weights with values drawn from a normal distribution.

Comments

Post a Comment