How to freeze model layers in PyTorch

Pretrained neural networks often contain a feature extractor that should stay fixed while a smaller task-specific head learns. In PyTorch, freezing a layer leaves it in the forward pass but stops autograd from storing gradients for that layer's parameters.

requires_grad controls whether a parameter participates in gradient calculation. Calling requires_grad_(False) on a module sets that flag on each parameter in the module, while assigning parameter.requires_grad = False works for narrower selections.

The optimizer should receive the parameters that are meant to move, and a backward pass plus one optimizer step makes the freeze boundary visible. A frozen layer should show requires_grad=False, no gradient tensor, and no weight change after optimizer.step().

1. Create a smoke script that freezes the first linear layer.

   layer_freeze_demo.py

   import torch
   from torch import nn
    
    
   torch.manual_seed(23)
    
   model = nn.Sequential(
       nn.Linear(4, 6),
       nn.ReLU(),
       nn.Linear(6, 2),
   )
    
   model[0].requires_grad_(False)
    
   trainable_parameters = [
       (name, parameter)
       for name, parameter in model.named_parameters()
       if parameter.requires_grad
   ]
   optimizer = torch.optim.SGD(
       [parameter for _, parameter in trainable_parameters],
       lr=0.1,
   )
   loss_fn = nn.CrossEntropyLoss()
    
   features = torch.tensor(
       [
           [0.5, -1.0, 0.3, 2.0],
           [1.0, 0.2, -0.4, 0.7],
           [-0.3, 1.2, 0.8, -1.1],
           [1.5, -0.7, 0.1, 0.4],
       ],
       dtype=torch.float32,
   )
   targets = torch.tensor([0, 1, 1, 0])
    
   before_step = {
       name: parameter.detach().clone()
       for name, parameter in model.named_parameters()
   }
    
   optimizer.zero_grad(set_to_none=True)
   loss = loss_fn(model(features), targets)
   loss.backward()
    
   print(f"torch_version={torch.__version__}")
   print("requires_grad:")
   for name, parameter in model.named_parameters():
       print(f"  {name}={parameter.requires_grad}")
    
   print(
       "optimizer_params="
       + ",".join(name for name, _ in trainable_parameters)
   )
   print(f"loss={loss.item():.4f}")
    
   print("gradient_state:")
   for name, parameter in model.named_parameters():
       if parameter.grad is None:
           print(f"  {name}=None")
       else:
           print(f"  {name}_norm={parameter.grad.norm().item():.4f}")
    
   optimizer.step()
    
   print("changed_after_step:")
   for name, parameter in model.named_parameters():
       changed = not torch.equal(before_step[name], parameter.detach())
       print(f"  {name}={changed}")

   The first Linear module is the frozen layer, and the second Linear module is the trainable head. The optimizer is built from parameters whose requires_grad flag is still True.

2. Run the smoke script.

   $ python layer_freeze_demo.py
   torch_version=2.12.1+cpu
   requires_grad:
     0.weight=False
     0.bias=False
     2.weight=True
     2.bias=True
   optimizer_params=2.weight,2.bias
   loss=0.6714
   gradient_state:
     0.weight=None
     0.bias=None
     2.weight_norm=0.2892
     2.bias_norm=0.1316
   changed_after_step:
     0.weight=False
     0.bias=False
     2.weight=True
     2.bias=True

3. Check the frozen and trainable parameter states in the output.

   0.weight and 0.bias show requires_grad=False, stay at None under gradient_state, and do not change after the optimizer step. 2.weight and 2.bias are the only optimizer parameters, receive gradients, and change after optimizer.step().

4. Freeze the intended submodule before constructing the real optimizer.

   model.backbone.requires_grad_(False)
    
   optimizer = torch.optim.Adam(
       (parameter for parameter in model.parameters() if parameter.requires_grad),
       lr=1e-4,
   )

   Freezing before optimizer creation keeps the trainable set explicit and avoids optimizer state for layers that should stay fixed. Move the model to its target device before constructing the optimizer when the training loop uses an accelerator.

5. Compare a frozen parameter before and after one training step when adapting the pattern to a project model.

   frozen_weight = next(model.backbone.parameters())
   frozen_before = frozen_weight.detach().clone()
    
   optimizer.zero_grad(set_to_none=True)
   loss = loss_fn(model(inputs), targets)
   loss.backward()
   optimizer.step()
    
   assert frozen_weight.grad is None
   assert torch.equal(frozen_before, frozen_weight.detach())

   The first assertion confirms the frozen parameter received no gradient. The second assertion confirms the optimizer step did not change that parameter.

6. Set the frozen submodule to evaluation mode when its BatchNorm or Dropout behavior should stay fixed.

   model.train()
   model.backbone.eval()

   requires_grad=False stops parameter gradients, but it does not change training/evaluation behavior. Leave the frozen submodule in training mode only when its buffers or stochastic layers should still behave like the rest of the training model.

7. Rebuild the optimizer when a frozen submodule is made trainable later.

   model.backbone.requires_grad_(True)
    
   optimizer = torch.optim.Adam(
       (parameter for parameter in model.parameters() if parameter.requires_grad),
       lr=1e-5,
   )

   If preserving existing optimizer state matters, add the newly trainable parameters with optimizer.add_param_group() instead of rebuilding the optimizer. Add a parameter group only for parameters that are not already present in the optimizer.

8. Remove the temporary smoke script.

   $ rm layer_freeze_demo.py