Pytorch
Deep Learning
Neural Networks
Machine Learning
Code Tutorial

How to iterate over layers 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

Iterating over layers is a basic PyTorch technique for inspection, freezing, replacement, and debugging. The important detail is that PyTorch offers several traversal APIs, and they answer slightly different questions: immediate children, recursive modules, named modules, and named parameters.

Start with children() and modules()

children() returns the immediate child modules of a model. modules() walks the entire nested tree and also includes the root module itself.

python
1import torch
2import torch.nn as nn
3
4model = nn.Sequential(
5    nn.Conv2d(3, 16, kernel_size=3, padding=1),
6    nn.ReLU(),
7    nn.Conv2d(16, 32, kernel_size=3, padding=1),
8    nn.ReLU(),
9    nn.AdaptiveAvgPool2d((1, 1)),
10    nn.Flatten(),
11    nn.Linear(32, 10),
12)
13
14print('children:')
15for layer in model.children():
16    print(type(layer).__name__)
17
18print('modules:')
19for layer in model.modules():
20    print(type(layer).__name__)

Use children() when you want top-level blocks. Use modules() when you want a recursive traversal.

Use Named Traversal for Stable References

Named iterators are more useful when you need logs, targeted changes, or optimizer grouping.

python
for name, layer in model.named_modules():
    print(name, '->', type(layer).__name__)

The names act like paths in the module hierarchy, which is much easier to reason about than printing only the class name.

Filter by Layer Type

A common task is applying an operation only to certain kinds of layers, such as freezing all convolutions.

python
1for name, layer in model.named_modules():
2    if isinstance(layer, nn.Conv2d):
3        for p in layer.parameters():
4            p.requires_grad = False
5        print('froze', name)

This pattern is simple and explicit, and it scales well to larger architectures.

Iterate Over Parameters When Optimization Matters

Sometimes the real target is not a layer object but the parameter tensors that belong to different parts of the network.

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

This is especially useful when building optimizer parameter groups.

python
1feature_params = []
2head_params = []
3
4for name, param in model.named_parameters():
5    if name.startswith('0.') or name.startswith('2.'):
6        feature_params.append(param)
7    else:
8        head_params.append(param)
9
10optimizer = torch.optim.Adam([
11    {'params': feature_params, 'lr': 1e-4},
12    {'params': head_params, 'lr': 1e-3},
13])

Use Hooks for Activation Inspection

If you need intermediate outputs rather than just the layer objects, register forward hooks.

python
1activations = {}
2
3
4def save_activation(name):
5    def hook(_module, _inp, out):
6        activations[name] = out.detach()
7    return hook
8
9handles = []
10for name, layer in model.named_modules():
11    if isinstance(layer, nn.Conv2d):
12        handles.append(layer.register_forward_hook(save_activation(name)))
13
14x = torch.randn(2, 3, 64, 64)
15_ = model(x)
16
17for name, value in activations.items():
18    print(name, value.shape)
19
20for h in handles:
21    h.remove()

Always remove hook handles when you are done. Otherwise they keep firing and can create confusing side effects.

Choose the Iterator That Matches the Task

A good rule is:

  • use children() for top-level structure
  • use named_modules() for recursive structural work
  • use named_parameters() for optimizer and freezing logic
  • use hooks when you need runtime activations

PyTorch gives you all of these because layer traversal is not one single problem.

When working with large pretrained architectures, it is often worth printing the traversal output once and saving that inspection in notes or tests. Many mistakes come from assuming a model's internal names or nesting structure without checking what PyTorch actually registered.

Common Pitfalls

A common mistake is using modules() when you expected only top-level layers. That can lead to repeated or overly broad modifications.

Another is forgetting that modules() includes the root model itself as the first item.

Developers also sometimes freeze parameters and then forget to rebuild the optimizer groups, which means the training setup no longer reflects the current requires_grad state.

Summary

  • Use children() for immediate child layers and modules() for recursive traversal.
  • Prefer named iterators when you need stable references.
  • Filter by type with isinstance for operations such as freezing or replacement.
  • Use named_parameters() when optimizer configuration depends on layer grouping.
  • Use hooks carefully when you need intermediate activations.

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.