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.
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.
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.
Use Kaiming, also called He initialization, when the layer is followed by a ReLU-like nonlinearity.
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.
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.
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.
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.
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:
- start with the module defaults or one standard initializer
- match the initializer to the activation function
- compare results under the same seed and optimizer settings
- 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
- How do I initialize weights in PyTorch?
- How do I keep track of the time the CPU is used vs the GPUs for deep learning?
- How do I know if tensorflow using cuda and cudnn or not?
- How do I load a keras saved model with custom Optimizer
- How do I load a local model with torch.hub.load?
- How do I load custom image based datasets into Pytorch for use with a CNN?
- How do I install tensorflow_text?
- How do I install TensorFlow's tensorboard?
.png&w=3840&q=75)
Tackling System Design Interview Problems
A short course that equips you with the skills to approach system design interviews methodically.
Start the free courseTrack 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.