PyTorch: nn.Module and registered parameters

pythontorchindependent-practice

Saved on this device

Full lesson · 4 learning steps Open any exercise for focused practice

Step 1 of 4 · intro card

PyTorch: nn.Module and registered parameters

Concept and contract

Make layers visible to optimizers, checkpoints and device moves.

An nn.Module groups trainable parameters and nested modules into a registered structure. Optimizers, state_dict and device moves discover that structure through the module APIs.

A plain Python list of layers can still execute a forward pass while hiding those layers from registration. Running successfully therefore does not establish that the intended parameters can be trained or saved.

Mechanism and worked example

TinyNet stores two linear layers in ModuleList and applies them in order.

TinyNet stores two linear layers in ModuleList and applies them in order. A two-feature batch becomes a one-output batch, and the helper counts every registered parameter, including biases.

The count is nine: six values in the first layer and three in the second. Explicit counts provide a useful sanity check before any optimizer or training loop is introduced.

Complete worked example

python

from bootstrap_lab import torch_setup
torch = torch_setup()

from torch import nn
class TinyNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.ModuleList([nn.Linear(2, 2), nn.Linear(2, 1)])
    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

def registered(batch_size):
    if type(batch_size) is not int or batch_size < 1:
        raise ValueError('positive batch size required')
    model = TinyNet()
    output = model(torch.ones(batch_size, 2))
    return {'shape': list(output.shape),
            'parameters': sum(parameter.numel() for parameter in model.parameters())}

Reasoning: TinyNet stores two linear layers in ModuleList and applies them in order. A two-feature batch becomes a one-output batch, and the helper counts every registered parameter, including biases.

Counterexamples and common mistakes

The counterexample replaces ModuleList with a regular list.

The counterexample replaces ModuleList with a regular list. The forward output retains a plausible shape, but model.parameters sees no nested layers.

That can produce an empty optimizer parameter list or incomplete checkpoints. Use ModuleList, ModuleDict, Sequential or named module attributes according to the intended structure, rather than depending on arbitrary Python containers to register children.

Counterexample: a runnable framework mistake

python

from bootstrap_lab import torch_setup
torch = torch_setup()

from torch import nn
class TinyNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = [nn.Linear(2, 2), nn.Linear(2, 1)]
    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

def registered(batch_size):
    if type(batch_size) is not int or batch_size < 1:
        raise ValueError('positive batch size required')
    model = TinyNet()
    output = model(torch.ones(batch_size, 2))
    return {'shape': list(output.shape),
            'parameters': sum(parameter.numel() for parameter in model.parameters())}

Reasoning: The counterexample replaces ModuleList with a regular list. The forward output retains a plausible shape, but model.parameters sees no nested layers.

Independent practice: one bounded framework change

Add one registered Linear layer mapping one output to one output after the existing pair.

Add one registered Linear layer mapping one output to one output after the existing pair. Preserve the batch-by-one output shape and update the expected parameter count to eleven.

This small architecture change reinforces registration and shape reasoning. Do not add an unrelated activation, change input feature count or substitute a hard-coded parameter total in the helper.

Step 1: Add a registered one-to-one final linear layer. Step 2: Extend the ModuleList rather than a plain list.

Practice starter: solve before reading the solution

python

from bootstrap_lab import torch_setup
torch = torch_setup()

from torch import nn
class TinyNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.ModuleList([nn.Linear(2, 2), nn.Linear(2, 1)])
    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

def registered(batch_size):
    # Remove this placeholder, then adapt the retained baseline:
    # Add a registered one-to-one final linear layer.
    raise NotImplementedError
    if type(batch_size) is not int or batch_size < 1:
        raise ValueError('positive batch size required')
    model = TinyNet()
    output = model(torch.ones(batch_size, 2))
    return {'shape': list(output.shape),
            'parameters': sum(parameter.numel() for parameter in model.parameters())}

Reasoning: Add a registered one-to-one final linear layer. Extend the ModuleList rather than a plain list.

Hint 1: Retain two input features and one final output.
Hint 2: Count parameters through model.parameters.

Solution, tests and trade-offs

The cases inspect batch sizes one and three and reject zero.

The cases inspect batch sizes one and three and reject zero. Practice checks the same output shapes with the larger observed parameter count.

No output values are asserted because random initialization is not the contract here. Counting registered parameters independently from the forward shape is what catches the deliberately hidden-layer counterexample.

Practice solution and acceptance checks

python

from bootstrap_lab import torch_setup
torch = torch_setup()

from torch import nn
class TinyNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.ModuleList([nn.Linear(2, 2), nn.Linear(2, 1), nn.Linear(1, 1)])
    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

def registered(batch_size):
    if type(batch_size) is not int or batch_size < 1:
        raise ValueError('positive batch size required')
    model = TinyNet()
    output = model(torch.ones(batch_size, 2))
    return {'shape': list(output.shape),
            'parameters': sum(parameter.numel() for parameter in model.parameters())}

Reasoning: The cases inspect batch sizes one and three and reject zero. Practice checks the same output shapes with the larger observed parameter count.

Interview reasoning and next step

As models grow, use named_parameters and state_dict keys to inspect ownership and checkpoint structure..

As models grow, use named_parameters and state_dict keys to inspect ownership and checkpoint structure.

Add nonlinearities deliberately when modeling nonlinear relationships; stacked linear layers alone remain an affine transformation. Keep architecture changes compatible with saved-state expectations, and treat a checkpoint shape mismatch as a real model-version issue rather than suppressing it without understanding missing or unexpected keys.

Optional stretch — after the core task works: Inspect named state_dict keys before and after an architecture change.

Computer practice: https://applydjinn.com/typorithm/python-frameworks-lab-20261105.zip. Install the pinned dependencies in an isolated environment, run verify.py --track torch, then edit the retained exercise baseline and check it with --practice pt-modules. The lab uses owned in-memory databases, local framework clients and tiny CPU tensors. It never reads the working ApplyDjinn database, contacts production feeds or downloads models.

Make layers visible to optimizers, checkpoints and device moves.