PyTorch modules and data

A PyTorch model is an nn.Module. It is a class that holds its weights and its sub-layers and defines what happens in forward. Training data reaches the model through a Dataset, which gives one example at a time, and a DataLoader, which groups them into shuffled batches.

Writing a module

Here is a two-layer network for 784 inputs, such as flattened 28×28 images:

import torch
import torch.nn as nn
import torch.nn.functional as F

class MLP(nn.Module):
    def __init__(self, d_in, d_hidden, d_out):
        super().__init__()
        self.fc1 = nn.Linear(d_in, d_hidden)
        self.fc2 = nn.Linear(d_hidden, d_out)

    def forward(self, x):
        return self.fc2(F.relu(self.fc1(x)))

model = MLP(784, 256, 10)
print(sum(p.numel() for p in model.parameters()))
print(list(model.state_dict().keys()))
print(model(torch.randn(3, 784)).shape)

Output:

203530
['fc1.weight', 'fc1.bias', 'fc2.weight', 'fc2.bias']
torch.Size([3, 10])

Check the count by hand. fc1 has 784 × 256 weights plus 256 biases, which is 200,960. fc2 has 256 × 10 weights plus 10 biases, which is 2,570. The total is 203,530.

Parameters and state_dict

• model.parameters() yields every trainable tensor. The optimizer receives this.
• model.state_dict() returns a dictionary of names and tensors, including buffers such as BatchNorm running statistics. Checkpoints store this.

The name of each weight comes from the attribute name, which is why fc1.weight appears.

class Bad(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = [nn.Linear(4, 4), nn.Linear(4, 4)]

class Good(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.ModuleList([nn.Linear(4, 4), nn.Linear(4, 4)])

print(len(list(Bad().parameters())), len(list(Good().parameters())))

Output:

0 4

The plain list reports zero parameters. The ModuleList reports four, two weights and two biases.

Dataset and DataLoader

A Dataset needs two methods: __len__ for the size and __getitem__ for one example. The DataLoader handles batching and shuffling.

from torch.utils.data import Dataset, DataLoader

class ToyDataset(Dataset):
    def __init__(self, n=1000, d=20):
        self.x = torch.randn(n, d)
        self.y = (self.x.sum(1) > 0).long()

    def __len__(self):
        return self.x.shape[0]

    def __getitem__(self, i):
        return self.x[i], self.y[i]

loader = DataLoader(ToyDataset(), batch_size=32, shuffle=True)
xb, yb = next(iter(loader))
print(xb.shape, yb.shape)
print(len(loader))

Output:

torch.Size([32, 20]) torch.Size([32])
32

1,000 examples in batches of 32 give 32 batches, because the last batch is smaller (1,000 = 31 × 32 + 8).

Common mistakes

• Storing layers in a plain list. Use nn.ModuleList or nn.Sequential so the parameters are found.
• Shuffling the validation set. Shuffle training data only. Validation results should come out in the same order every run.
• Forgetting drop_last matters. A small final batch can upset BatchNorm. Set drop_last=True if that matters for your model.

Check yourself

1. Why does state_dict() include BatchNorm running statistics when they are not parameters? Answer: they are buffers that the model needs at inference time, so a checkpoint must keep them.
2. A Dataset has 500 examples and batch_size=64. How many batches per epoch? Answer: 8, because 500 / 64 = 7.8 and the last batch is partial.

Key takeaways

• An nn.Module owns its parameters and defines forward.
• parameters() feeds the optimizer. state_dict() is what you save.
• Register sub-layers so PyTorch can see them. Use ModuleList, not a list.
• Dataset yields one example. DataLoader batches and shuffles.

Further learning

Next, put these pieces into a training loop in the canonical training loop.