PyTorch
weight initialization
deep learning
machine learning
neural networks

How do I initialize weights in PyTorch?

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

Weight initialization sets the starting point for optimization in a neural network. In PyTorch, you usually do not need to invent a strategy from scratch, but you do need to choose an initializer that matches the layer type and activation function if you want predictable training behavior.

PyTorch Modules Already Initialize Parameters

A useful first fact is that nn.Linear, nn.Conv2d, and other standard modules already come with default parameter initialization. That means a model is not broken simply because you did not manually initialize every tensor.

python
1import torch
2import torch.nn as nn
3
4layer = nn.Linear(8, 4)
5print(layer.weight)
6print(layer.bias)

The defaults are often good enough for baseline experiments. Custom initialization becomes more useful when:

  • training is unstable
  • you are reproducing a published architecture
  • you want explicit control for experiments
  • you added custom layers with no sensible default

Match the Initializer to the Activation

The two most common initialization families are Xavier and Kaiming.

Use Xavier, also called Glorot, when the layer is followed by symmetric activations such as tanh.

python
1import torch.nn as nn
2
3layer = nn.Linear(32, 64)
4nn.init.xavier_uniform_(layer.weight)
5nn.init.zeros_(layer.bias)

Use Kaiming, also called He initialization, when the layer is followed by a ReLU-like nonlinearity.

python
1import torch.nn as nn
2
3layer = nn.Linear(32, 64)
4nn.init.kaiming_uniform_(layer.weight, nonlinearity="relu")
5nn.init.zeros_(layer.bias)

These methods exist because different activations preserve variance differently as signals move through the network.

Initialize a Full Model with apply

The standard PyTorch pattern is to write one function and apply it to the whole model.

python
1import torch
2import torch.nn as nn
3
4class Net(nn.Module):
5    def __init__(self):
6        super().__init__()
7        self.fc1 = nn.Linear(16, 32)
8        self.fc2 = nn.Linear(32, 10)
9
10    def forward(self, x):
11        x = torch.relu(self.fc1(x))
12        return self.fc2(x)
13
14
15def init_weights(module):
16    if isinstance(module, nn.Linear):
17        nn.init.kaiming_normal_(module.weight, nonlinearity="relu")
18        nn.init.zeros_(module.bias)
19
20
21model = Net()
22model.apply(init_weights)

This is cleaner than initializing each layer manually after construction, especially as models get deeper.

Convolutional Layers Follow the Same Idea

For convolutional networks, the same activation logic applies. ReLU-heavy conv stacks often use Kaiming initialization.

python
1import torch.nn as nn
2
3
4def init_conv(module):
5    if isinstance(module, nn.Conv2d):
6        nn.init.kaiming_normal_(module.weight, nonlinearity="relu")
7        if module.bias is not None:
8            nn.init.zeros_(module.bias)

You can combine rules for both linear and convolutional layers inside one initializer function if the model uses both.

Bias Initialization Is Usually Simpler

In many networks, zero bias initialization is perfectly acceptable.

python
nn.init.zeros_(layer.bias)

Biases rarely need the same level of experimentation as weights. The main goal is usually consistency and simplicity unless a specific architecture paper says otherwise.

Reproducibility Matters During Comparison

If you are comparing two initialization strategies, fix the random seed. Otherwise you are mixing the effect of initialization choice with the effect of random sampling.

python
import torch

torch.manual_seed(42)

With the seed fixed, you can compare Xavier against Kaiming or defaults with less noise in the result.

Initialization Does Not Replace Good Training Setup

Poor initialization can hurt training, but it is not the only cause of bad results. Learning rate, normalization, batch size, optimizer choice, and data preprocessing can all dominate the outcome.

A practical workflow is:

  1. start with the module defaults or one standard initializer
  2. match the initializer to the activation function
  3. compare results under the same seed and optimizer settings
  4. only then treat initialization as a tuning variable

This keeps initialization in perspective instead of turning it into a superstition.

Common Built-in Options

PyTorch exposes several useful initializers in torch.nn.init, including:

  • 'xavier_uniform_'
  • 'xavier_normal_'
  • 'kaiming_uniform_'
  • 'kaiming_normal_'
  • 'normal_'
  • 'uniform_'
  • 'zeros_'
  • 'ones_'

That means most common initialization needs are already covered by the library. You rarely need to fill tensors by hand.

Common Pitfalls

A common mistake is reinitializing weights after training has already begun. Initialization should happen once, before optimization starts.

Another mistake is choosing Xavier for a strongly ReLU-based network just because it is well known. Kaiming often matches ReLU stacks better.

Developers also sometimes compare initialization strategies without fixing the seed or keeping training settings constant, which makes the comparison unreliable.

Finally, if the model is still unstable after sensible initialization, the problem may be the optimizer, learning rate, or data pipeline rather than the initializer itself.

Summary

  • PyTorch layers already have default initialization, so manual initialization is optional.
  • Use Xavier for tanh-like activations and Kaiming for ReLU-like activations.
  • Apply initialization cleanly with model.apply(...).
  • Zero bias initialization is a common reasonable default.
  • Evaluate initialization as part of the whole training setup, not in isolation.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.