Freezing layers in PyTorch

Introduction

Freezing layers helps keep some parts of a model fixed so they don't change during training. This saves time and keeps learned knowledge safe.

When using a pre-trained model and you want to keep early layers unchanged.

When training a model on a small dataset to avoid overfitting.

When you want to speed up training by updating only some layers.

When transferring learning from one task to another similar task.

When experimenting with which parts of a model to train or keep fixed.

Syntax

PyTorch

for param in model.layer.parameters():
    param.requires_grad = False

Setting requires_grad to False tells PyTorch not to update these parameters during training.

You can freeze entire layers or just parts by selecting the right parameters.

Examples

Freeze all parameters in the first convolutional layer.

PyTorch

for param in model.conv1.parameters():
    param.requires_grad = False

Freeze all layers except those with 'fc' in their name (usually fully connected layers).

PyTorch

for name, param in model.named_parameters():
    if 'fc' not in name:
        param.requires_grad = False

Freeze the whole model except the final fully connected layer.

PyTorch

for param in model.parameters():
    param.requires_grad = False

for param in model.fc.parameters():
    param.requires_grad = True

Sample Model

This code creates a simple model with two layers. It freezes the first layer so it does not update during training. The optimizer only updates the second layer. After one training step, we print gradients to see which layers changed.

PyTorch

import torch
import torch.nn as nn
import torch.optim as optim

# Simple model with two linear layers
class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(5, 3)
        self.fc2 = nn.Linear(3, 1)
    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

model = SimpleModel()

# Freeze first layer
for param in model.fc1.parameters():
    param.requires_grad = False

# Optimizer only updates parameters with requires_grad=True
optimizer = optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=0.1)

# Dummy data
inputs = torch.randn(4, 5)
targets = torch.randn(4, 1)

criterion = nn.MSELoss()

# Training step
model.train()
outputs = model(inputs)
loss = criterion(outputs, targets)
loss.backward()
optimizer.step()

# Check which parameters were updated
fc1_grad = model.fc1.weight.grad
fc2_grad = model.fc2.weight.grad

print("fc1 weight grad:", fc1_grad)
print("fc2 weight grad:", fc2_grad)

Important Notes

Frozen layers do not get gradients, so their weights stay the same.

Always filter parameters with requires_grad=True when creating the optimizer.

You can unfreeze layers later by setting requires_grad back to True.

Summary

Freezing layers stops them from updating during training.

It is useful for transfer learning and saving training time.

Set requires_grad=False on parameters to freeze them.