Lesson 3

Python functions to reusable neural layers

Introduction

A linear rule can turn two input features, 2 and 3, into three class scores, 2, 3, and 5. A PyTorch module packages this reusable function and registers the weights that training must update.

Learning goal

Build a small module, trace its input and output shapes, and count its registered parameters.

Before you start

Python functions, tensors, matrix multiplication, classes as named containers, and basic model scores.

Lesson plan

  1. Start from a plain function and learn only the class syntax needed for modules.
  2. Run a fixed linear example and interpret logits, predictions, shapes, and parameter count.
  3. Inspect registered parameters and explain why stored layers must be module attributes.

Start with the function

A model is first a function: input numbers go in, and output numbers come out. For example, a linear function can calculate x @ W + b. Here W is a table of weights, b is a bias vector, and @ is matrix multiplication. During training, both W and b should be found and updated automatically.

nn.Module is PyTorch's standard container for such a function. It keeps track of nested layers and trainable parameters. This makes model.parameters(), device movement, saving, loading, training mode, and evaluation mode work consistently.

Just enough Python classes

ClassA recipe for making objects that hold data and functions together.
ObjectOne value made from a class, such as one particular model.
__init__The setup function called when the object is created.
selfThe current object. self.layer saves a layer inside that object.
forwardThe function that transforms model input into output.
ParameterA tensor registered for learning, usually created inside a PyTorch layer.

super().__init__() runs the setup code from nn.Module. It must run before assigning submodules. We call the object as model(x), not usually as model.forward(x). The call syntax lets PyTorch run hooks and other module machinery around forward.

Worked example: scores for three classes

Our input contains two measured features per example. The model returns three raw class scores, called logits. A logit is not a probability; it can be any real number. Shape moves from (batch, 2) to (batch, 3). We set fixed weights only to make the expected output easy to inspect.

Predict first: the third output row of weights is [1, 1]. For input [2, 3], predict that logit before reading the output. Also predict the total parameter count: three classes need two weights and one bias each.

import torch
from torch import nn

class TinyClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.output = nn.Linear(in_features=2, out_features=3)

    def forward(self, x):
        return self.output(x)

model = TinyClassifier()
with torch.no_grad():
    model.output.weight.copy_(torch.tensor([
        [1.0, 0.0],
        [0.0, 1.0],
        [1.0, 1.0],
    ]))
    model.output.bias.zero_()

x = torch.tensor([[2.0, 3.0],
                  [1.0, -1.0]])
logits = model(x)
print(logits)
print(tuple(logits.shape))

# Expected values:
# [[2.0, 3.0, 5.0], [1.0, -1.0, 0.0]]
# (2, 3)

For the first row, class 0 receives 2, class 1 receives 3, and class 2 receives 2 + 3 = 5. The largest score is class 2. The second row predicts class 0. torch.argmax(logits, dim=1) therefore returns [2, 0].

Registered parameters

Assigning nn.Linear to self.output registers its weight and bias with the outer model. PyTorch can now find them recursively. This is why layers should be stored as module attributes, or inside containers such as nn.Sequential and nn.ModuleList. An ordinary Python list of layers is not registered in the same way.

for name, parameter in model.named_parameters():
    print(name, tuple(parameter.shape))

# Expected:
# output.weight (3, 2)
# output.bias (3,)

The weight shape is (out_features, in_features). Each of the three output classes has two weights. The bias has one value per output. The batch axis is not part of the parameter shape. Because these registered parameters have requires_grad=True, autograd can calculate gradients for them.

The prediction is 5. The parameter count is 3 × 2 + 3 = 9. A common wrong count includes the batch size. Parameters describe the reusable rule, so changing from two examples to twenty examples does not add parameters.

Common pitfalls

  • Call super().__init__() before creating child layers.
  • Define layers in __init__, not inside forward; otherwise new weights may be created on every call.
  • Pass floating-point features to nn.Linear. Integer token IDs first need an embedding or another conversion.
  • Do not apply softmax before nn.CrossEntropyLoss; that loss expects raw logits.
  • Keep input and parameters on the same device and use compatible dtypes.

Try it

Change the model to accept four input features and produce two scores. For a batch of five examples, what should the output shape be? How many trainable scalar values does the linear layer contain?

Reveal the worked answer 
import torch
from torch import nn

layer = nn.Linear(4, 2)
x = torch.zeros(5, 4)
y = layer(x)
print(tuple(y.shape))
print(sum(p.numel() for p in layer.parameters()))

# Expected:
# (5, 2)
# 10

The weight contains 2 × 4 = 8 values. The bias adds two more, for ten total trainable values.

Recap

A model is a function, and nn.Module is PyTorch's organized home for that function. __init__ creates and stores layers; forward describes the data flow. self means the current model object. Registered layers expose their parameters to optimizers, saving tools, and device movement. Always reason about input and output shapes before training.

Reference: PyTorch nn.Module documentation.