PyTorch
freeze weights
update param_groups
machine learning
deep learning

pytorch freeze weights and update param_groups

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

Freezing layers in PyTorch means telling autograd not to compute gradients for selected parameters. Updating optimizer parameter groups is the next step when you want different learning rates or when you later unfreeze part of the model. The important detail is that changing requires_grad and changing optimizer groups are related, but they are not the same operation.

Freeze Parameters with requires_grad

The core freezing step is simple:

python
1import torch
2import torch.nn as nn
3
4model = nn.Sequential(
5    nn.Linear(10, 20),
6    nn.ReLU(),
7    nn.Linear(20, 2),
8)
9
10for param in model[0].parameters():
11    param.requires_grad = False

Now the first linear layer will not receive gradients during backpropagation.

Build the Optimizer from Trainable Parameters

If a layer is frozen from the start, the cleanest approach is to pass only trainable parameters to the optimizer.

python
1optimizer = torch.optim.Adam(
2    [p for p in model.parameters() if p.requires_grad],
3    lr=1e-3
4)

This avoids carrying frozen parameters in the optimizer state unnecessarily.

Use Parameter Groups for Different Learning Rates

Parameter groups are useful when different parts of the model should train at different speeds.

python
1optimizer = torch.optim.Adam([
2    {"params": model[0].parameters(), "lr": 1e-4},
3    {"params": model[2].parameters(), "lr": 1e-3},
4])

This is common in transfer learning, where a pretrained backbone gets a smaller learning rate than a newly added classifier head.

Unfreeze Later and Update the Optimizer

If you unfreeze layers after some epochs, you must also make sure the optimizer knows about the newly trainable parameters. One option is to rebuild the optimizer. Another is to add a new parameter group.

python
1for param in model[0].parameters():
2    param.requires_grad = True
3
4optimizer.add_param_group({
5    "params": model[0].parameters(),
6    "lr": 1e-4
7})

If those parameters were never part of the optimizer before, add_param_group is appropriate.

If they were already present but frozen, rebuilding the optimizer is often simpler and less error-prone, especially when optimizer state should be reset consistently.

Verify What Will Actually Train

Do not assume freezing worked just because you set a flag once. Check trainable parameter counts or names before training.

python
for name, param in model.named_parameters():
    print(name, param.requires_grad)

This is especially important in larger models where one loop may miss nested modules or repeated blocks.

Fine-Tuning Usually Happens in Stages

A common workflow is to train a new head first, then unfreeze some or all of the backbone later. Parameter groups are valuable here because the second stage often needs a lower learning rate for pretrained layers than for the head.

That staged plan should be visible in code. If the optimizer, requires_grad flags, and learning rates do not all change together, the fine-tuning schedule is usually incomplete.

Know the Difference Between Gradient Flow and Optimizer State

A frozen parameter with requires_grad = False will not accumulate gradients. But if the parameter is still in the optimizer, the optimizer object still holds references and potentially state for it. That is why optimizer configuration should match the training plan rather than being treated as an unrelated detail.

Common Pitfalls

  • Setting requires_grad = False but forgetting to align optimizer parameter groups with the new training plan.
  • Assuming add_param_group is always better than rebuilding the optimizer after unfreezing.
  • Freezing a module after the optimizer is created and never checking what the optimizer still contains.
  • Giving the same learning rate to pretrained and newly initialized layers without thinking about stability.
  • Forgetting to verify trainable parameters before launching a long training run.

Summary

  • Freeze PyTorch weights by setting requires_grad = False on the relevant parameters.
  • Build the optimizer from trainable parameters when freezing from the start.
  • Use parameter groups when different layers need different learning rates.
  • If you unfreeze later, update the optimizer deliberately with a new group or a rebuild.
  • Stage fine-tuning carefully so optimizer state and trainable flags stay aligned.

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.