Building with nn.Module

Every network in PyTorch, from a tiny digit reader to the largest language models, has the same simple shape. It's a Python class with two parts. A list of its layers, and a recipe for how data flows through them. Once you've seen this pattern, you can read almost any model in the world.

• Every network has the same shape
• Layers, then a recipe
• Read any model in the world

A network as a class

We make a class that inherits from n n dot module. In init, we create the layers. Flatten turns an image into a long list. Linear is a layer of neurons, connecting every input to every output. Then forward is the recipe. Flatten, pass through the hidden layer and an activation, and then the output. Run it. PyTorch prints a clean summary. Our little network has over a hundred thousand weights. And the output shape is four rows by ten, which is one score for each digit, for each image in the batch.

import torch
from torch import nn

class DigitNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.hidden = nn.Linear(28 * 28, 128)
        self.relu = nn.ReLU()
        self.output = nn.Linear(128, 10)

    def forward(self, x):
        x = self.flatten(x)
        x = self.relu(self.hidden(x))
        return self.output(x)

model = DigitNet()
print(model)
print("Weights to learn:", sum(p.numel() for p in model.parameters()))
print("Output shape:", model(torch.rand(4, 1, 28, 28)).shape)

DigitNet(
  (flatten): Flatten(start_dim=1, end_dim=-1)
  (hidden): Linear(in_features=784, out_features=128, bias=True)
  (relu): ReLU()
  (output): Linear(in_features=128, out_features=10, bias=True)
)
Weights to learn: 101770
Output shape: torch.Size([4, 10])

The language of layers

Three words will keep coming up. A linear layer multiplies its inputs by weights. An activation adds a bend, so that the network can learn curves. And the raw outputs at the end are called logits, one score per class. We don't convert these into chances ourselves, because the loss function handles that, in a more stable way. Understanding these three words will let you read most model code.

Summary