Skip to content

← nn.Module Mechanics step 7 of 7

Medium Primitives

Initialise the whole tree

Jump to the task

Use Module.apply when initialization depends on the kind of layer. The callback receives each registered child module recursively, followed by its parent, including the root module itself:

import torch
import torch.nn as nn

def init_linear(module):
    if isinstance(module, nn.Linear):
        nn.init.ones_(module.weight)
        if module.bias is not None:
            nn.init.zeros_(module.bias)

model = nn.Sequential(nn.Linear(2, 2), nn.Sequential(nn.Linear(2, 2)))
model.apply(init_linear)
print(model(torch.ones(2)).tolist())  # [4.0, 4.0]

The callback ignores the Sequential containers and initializes both Linear layers. The bias check also lets it handle a Linear created with bias=False. As with parameters() and to(), the children must be registered for apply to find them. Modules kept only in a plain Python list are not reached.

Give weights and biases different treatment

Setting every parameter to one would also set every bias to one. With a module callback, you can access weight and bias separately and choose an initialization suited to that layer.

There is no single initialization recipe for every architecture. All-ones weights make this exercise easy to check, but they are not a general training recommendation: identical weights can preserve unwanted symmetry between units. A loop using named_parameters() is another option when parameter names give you enough information; apply is convenient when you need the layer type.

Initialization and gradient recording

The nn.init functions perform their updates without recording gradients. A direct in-place update of a trainable leaf normally needs your own no-grad context:

layer = nn.Linear(2, 2)
with torch.no_grad():
    layer.weight.fill_(1.0)
nn.init.zeros_(layer.bias)
print(layer(torch.ones(2)).tolist())  # [2.0, 2.0]

Initialize before computing the forward pass that you plan to differentiate. No-grad avoids recording the update; it does not make changing a tensor safe if an existing backward calculation still needs its old value.

Count tensors or count elements

len(list(model.parameters())) counts parameter tensors. sum(p.numel() for p in model.parameters()) counts their elements, including frozen parameters. Add if p.requires_grad to count only trainable elements. The iterator removes duplicate references to the same Parameter object.

Your task

def init_and_count(depth: int) -> dict

Build a Sequential containing depth independent Linear(2, 2) layers. Use apply to set every Linear weight to ones and every bias to zeros, then run the model on torch.ones(2). Return:

  • param_total: the sum of all parameter values, rounded to four decimals.
  • n_elements: the total number of parameter elements.
  • output: the output as a list, with each value rounded to four decimals.

Each layer contributes four weight elements and two bias elements. It doubles each coordinate when both input coordinates are equal. For depth=2, return {"param_total": 8.0, "n_elements": 12, "output": [4.0, 4.0]}.