Building Models with nn.Module

Week 2 of 18 · Foundations · 7 days

Week 02 · Foundations

Building Models with nn.Module

Week 02 · Day 1 of 7

nn.Module

Custom layers, registered parameters, and the list that silently does not train

nn.Sequential is fine until the moment your model does anything other than pass one thing straight through, which is almost immediately. A residual connection, two inputs, a branch that rejoins later: none of those fit in a list. nn.Module does.

The smallest possible module

import torch
from torch import nn

class MyLinear(nn.Module):
    def __init__(self, n_in, n_out):
        super().__init__()                  # never forget this line
        self.weight = nn.Parameter(torch.randn(n_out, n_in) * 0.1)
        self.bias = nn.Parameter(torch.zeros(n_out))

    def forward(self, x):
        return x @ self.weight.T + self.bias

layer = MyLinear(4, 3)
x = torch.randn(2, 4)
print('output shape', tuple(layer(x).shape))
print('\nparameters it registered:')
for name, prm in layer.named_parameters():
    print('  %-8s %s' % (name, tuple(prm.shape)))

output shape (2, 3)

parameters it registered:
  weight   (3, 4)
  bias     (3,)

What nn.Parameter is for

A plain tensor assigned to self is invisible. Wrap it in nn.Parameter and the module registers it, which means it appears in model.parameters(), gets passed to the optimiser, moves with .to(device) and is saved in state_dict(). Forget the wrapper and the tensor simply never trains, silently, while everything else does.

import torch
from torch import nn

class Forgetful(nn.Module):
    def __init__(self):
        super().__init__()
        self.good = nn.Parameter(torch.zeros(3))
        self.bad = torch.zeros(3)               # not a Parameter

m = Forgetful()
print('registered parameters:', [n for n, _ in m.named_parameters()])
print('the optimiser would update %d tensor(s)'
      % len(list(m.parameters())))
print('\nself.bad exists and is used in forward, but never learns.')

registered parameters: ['good']
the optimiser would update 1 tensor(s)

self.bad exists and is used in forward, but never learns.

Modules contain modules

import torch
from torch import nn

class Block(nn.Module):
    def __init__(self, size):
        super().__init__()
        self.fc = nn.Linear(size, size)
        self.act = nn.ReLU()

    def forward(self, x):
        return self.act(self.fc(x))

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.stem = nn.Linear(784, 64)
        self.blocks = nn.ModuleList([Block(64) for _ in range(3)])
        self.head = nn.Linear(64, 10)

    def forward(self, x):
        x = torch.relu(self.stem(x.flatten(1)))
        for block in self.blocks:
            x = block(x)
        return self.head(x)

net = Net()
print('parameters %d' % sum(p.numel() for p in net.parameters()))
print('output', tuple(net(torch.randn(5, 1, 28, 28)).shape))
print('\nnamed modules:')
for name, mod in net.named_children():
    print('  %-8s %s' % (name, type(mod).__name__))

parameters 63370
output (5, 10)

named modules:
  stem     Linear
  blocks   ModuleList
  head     Linear

A plain Python list hides its modules

self.blocks = [Block(64) for _ in range(3)] looks identical and is broken in the same way as a plain tensor: the modules in it are not registered, so their parameters do not reach the optimiser and do not move to the GPU. Use nn.ModuleList when you need a list and nn.ModuleDict when you need a dictionary. The symptom is a model that trains, slowly, using only the layers you happened to assign directly.

A residual connection, which Sequential cannot express

import torch
from torch import nn

class Residual(nn.Module):
    def __init__(self, size):
        super().__init__()
        self.fc1 = nn.Linear(size, size)
        self.fc2 = nn.Linear(size, size)

    def forward(self, x):
        h = torch.relu(self.fc1(x))
        h = self.fc2(h)
        return torch.relu(x + h)        # the input rejoins the output

block = Residual(16)
x = torch.randn(4, 16)
print('in ', tuple(x.shape), ' out', tuple(block(x).shape))
print('\nthat one addition is the idea week 6 is built on.')

in  (4, 16)  out (4, 16)

that one addition is the idea week 6 is built on.

Day 1 takeaway

Subclass nn.Module, call super().__init__() first, wrap learnable tensors in nn.Parameter and lists of layers in nn.ModuleList. Anything that is not registered does not train, and nothing warns you.

Week 02 · Day 2 of 7

Losses

Logits not probabilities, stability at the extremes, and weighting

By Pawan Rama Mali793 words

The loss function is where you tell the model what counts as wrong. Choosing the wrong one is not a tuning mistake, it is a specification mistake, and the model will optimise exactly what you asked for.

The three you will use most

Task	Final layer	Loss	Target format
Binary classification	1 unit, no activation	`BCEWithLogitsLoss`	float 0.0 or 1.0
Multi-class, one label	C units, no activation	`CrossEntropyLoss`	integer class index
Multi-label	C units, no activation	`BCEWithLogitsLoss`	float vector of 0s and 1s
Regression	1 unit, no activation	`MSELoss` or `L1Loss`	float
Regression with outliers	1 unit, no activation	`HuberLoss`	float

Notice that no row has a softmax or a sigmoid in it

CrossEntropyLoss applies log_softmax internally, and BCEWithLogitsLoss applies the sigmoid. Doing it yourself first is not a harmless duplication: it makes the gradients smaller and the arithmetic less stable, and it trains a worse model without raising anything. The WithLogits in the name is the library telling you it has already been handled.

import torch
from torch import nn

logits = torch.tensor([[2.0, 0.5, -1.0]])
target = torch.tensor([0])

correct = nn.CrossEntropyLoss()(logits, target)

# the same thing done by hand, to show what it contains
log_probs = logits - logits.exp().sum().log()
by_hand = -log_probs[0, target[0]]

print('CrossEntropyLoss %.6f' % correct.item())
print('by hand          %.6f' % by_hand.item())

# and the mistake
wrong = nn.CrossEntropyLoss()(torch.softmax(logits, dim=1), target)
print('\nwith a softmax applied first: %.6f' % wrong.item())
print('different, smaller gradients, and no error raised.')

CrossEntropyLoss 0.241311
by hand          0.241311

with a softmax applied first: 0.701717
different, smaller gradients, and no error raised.

Numerical stability, demonstrated

import torch
from torch import nn

logits = torch.tensor([[-30.0]], requires_grad=True)
target = torch.tensor([[1.0]])

stable = nn.BCEWithLogitsLoss()(logits, target)
print('BCEWithLogitsLoss   %.6f' % stable.item())

probs = torch.sigmoid(logits)
naive = nn.BCELoss()(probs, target)
print('sigmoid then BCELoss %.6f' % naive.item())
print('agree here, at a probability of %.3e' % probs.item())

extreme = torch.tensor([[-200.0]], requires_grad=True)
print('\nnow at -200, where the sigmoid underflows to exactly 0:')
print('  with logits %.4f' % nn.BCEWithLogitsLoss()(extreme, target).item())
print('  the naive way %.4f'
      % nn.BCELoss()(torch.sigmoid(extreme), target).item())
print('\n100 is BCELoss clamping. The true loss is 200, so the')
print('gradient is wrong exactly where the model is most wrong.')

BCEWithLogitsLoss   30.000000
sigmoid then BCELoss 30.000000
agree here, at a probability of 9.358e-14

now at -200, where the sigmoid underflows to exactly 0:
  with logits 200.0000
  the naive way 100.0000

100 is BCELoss clamping. The true loss is 200, so the
gradient is wrong exactly where the model is most wrong.

At a logit of -30 the two agree perfectly, which is worth noticing: this is not a problem you will see in a toy example. It appears when a model becomes confident, which is to say late in training on the examples it has learned best, and it appears as a loss that stops decreasing for no visible reason.

Weighting classes

import torch
from torch import nn

# Ten classes where class 0 is nine times more common than the rest.
torch.manual_seed(0)
counts = torch.tensor([900.0] + [100.0] * 9)
weights = counts.sum() / (len(counts) * counts)
print('class counts ', [int(c) for c in counts.tolist()])
print('loss weights ', [round(w, 3) for w in weights.tolist()])

logits = torch.randn(4, 10)
target = torch.tensor([0, 0, 0, 7])
print('\nunweighted loss %.4f' % nn.CrossEntropyLoss()(logits, target).item())
print('weighted loss   %.4f'
      % nn.CrossEntropyLoss(weight=weights)(logits, target).item())

class counts  [900, 100, 100, 100, 100, 100, 100, 100, 100, 100]
loss weights  [0.2, 1.8, 1.8, 1.8, 1.8, 1.8, 1.8, 1.8, 1.8, 1.8]

unweighted loss 2.7360
weighted loss   3.0033

Weighting makes a mistake on a rare class cost more, which is one of three ways to handle imbalance. The others are resampling the data and leaving the loss alone but moving the decision threshold afterwards. The third is usually the best and the least used, because it does not require retraining anything.

Regression losses are not interchangeable

import torch
from torch import nn

pred = torch.zeros(5)
target = torch.tensor([0.1, -0.2, 0.05, 0.0, 8.0])   # one outlier

print('%-14s %10s' % ('loss', 'value'))
for name, fn in [('MSELoss', nn.MSELoss()),
                 ('L1Loss', nn.L1Loss()),
                 ('HuberLoss', nn.HuberLoss(delta=1.0))]:
    print('%-14s %10.4f' % (name, fn(pred, target).item()))

print('\ngradient contributed by the outlier alone:')
for name, fn in [('MSELoss', nn.MSELoss()),
                 ('L1Loss', nn.L1Loss()),
                 ('HuberLoss', nn.HuberLoss(delta=1.0))]:
    p_ = torch.zeros(5, requires_grad=True)
    fn(p_, target).backward()
    print('  %-12s %.4f' % (name, p_.grad[4].item()))

loss                value
MSELoss           12.8105
L1Loss             1.6700
HuberLoss          1.5052

gradient contributed by the outlier alone:
  MSELoss      -3.2000
  L1Loss       -0.2000
  HuberLoss    -0.2000

Day 2 takeaway

Pass logits, not probabilities. CrossEntropyLoss takes integer class indices and does the softmax itself; BCEWithLogitsLoss does the sigmoid itself and stays stable at extreme values. Squared error lets one outlier dominate the gradient, which is a choice you should make on purpose.

Week 02 · Day 3 of 7

Splits and Early Stopping

Watching a network memorise, and keeping the weights that were best

By Pawan Rama Mali1277 words

A model that scores well on the data it was fitted to has told you nothing. This is the same discipline as any other machine learning, with one addition: a deep network can memorise a training set completely, so the gap between the two numbers is usually larger and always more important.

Three splits, not two

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf)

# Subsets, so every page finishes while you watch it. The lessons are the
# same at full size and the numbers are a little higher.
train_set = Subset(train_all, range(8000))
val_set = Subset(train_all, range(8000, 10000))
test_set = Subset(test_all, range(2000))
print('train %d   validation %d   test %d'
      % (len(train_set), len(val_set), len(test_set)))
print()
print('train      : gradients are computed from it')
print('validation : you look at it after every epoch and make decisions')
print('test       : you look at it once, at the end')

train 8000   validation 2000   test 2000

train      : gradients are computed from it
validation : you look at it after every epoch and make decisions
test       : you look at it once, at the end

The validation set is training data too, in the way that matters

Every time you stop training because validation loss rose, pick an architecture because it validated better, or tune a learning rate against it, you have used the validation set to make a choice. Do that fifty times and the validation score is optimistic in exactly the way a training score is. That is what the test set is held back for, and why it is worth almost nothing after you have looked at it twice.

Watching a model memorise

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf)

# Subsets, so every page finishes while you watch it. The lessons are the
# same at full size and the numbers are a little higher.
train_set = Subset(train_all, range(8000))
val_set = Subset(train_all, range(8000, 10000))
test_set = Subset(test_all, range(2000))
def fit(model, epochs=4, lr=0.1, train=None, val=None, log=True):
    """The five steps from week 1, wrapped up so the pages stay short."""
    train = train if train is not None else DataLoader(train_set,
                                                       batch_size=64,
                                                       shuffle=True)
    val = val if val is not None else DataLoader(val_set, batch_size=256)
    loss_fn = nn.CrossEntropyLoss()
    opt = torch.optim.SGD(model.parameters(), lr=lr)
    history = []
    for epoch in range(1, epochs + 1):
        model.train()
        seen = total = 0
        for xb, yb in train:
            opt.zero_grad()
            loss = loss_fn(model(xb), yb)
            loss.backward()
            opt.step()
            total += loss.item() * yb.numel()
            seen += yb.numel()
        model.eval()
        right = count = 0
        with torch.no_grad():
            for xb, yb in val:
                right += (model(xb).argmax(1) == yb).sum().item()
                count += yb.numel()
        history.append((total / seen, right / count))
        if log:
            print('epoch %d  train loss %.4f  val accuracy %.4f'
                  % (epoch, total / seen, right / count))
    return history
torch.manual_seed(0)

# A deliberately oversized model on a deliberately tiny training set.
small = Subset(train_all, range(300))
loader = DataLoader(small, batch_size=32, shuffle=True)

model = nn.Sequential(nn.Flatten(), nn.Linear(784, 512), nn.ReLU(),
                      nn.Linear(512, 512), nn.ReLU(), nn.Linear(512, 10))
print('parameters %d for %d training images\n'
      % (sum(p.numel() for p in model.parameters()), len(small)))

loss_fn = nn.CrossEntropyLoss()
opt = torch.optim.SGD(model.parameters(), lr=0.1)
val_loader = DataLoader(val_set, batch_size=256)

print('%6s %14s %14s %14s' % ('epoch', 'train loss', 'train acc', 'val acc'))
for epoch in range(1, 41):
    model.train()
    for xb, yb in loader:
        opt.zero_grad()
        loss_fn(model(xb), yb).backward()
        opt.step()
    if epoch % 8 == 0 or epoch == 1:
        model.eval()
        with torch.no_grad():
            tr_right = sum((model(xb).argmax(1) == yb).sum().item()
                           for xb, yb in loader)
            tr_loss = sum(loss_fn(model(xb), yb).item() * yb.numel()
                          for xb, yb in loader) / len(small)
            va_right = sum((model(xb).argmax(1) == yb).sum().item()
                           for xb, yb in val_loader)
        print('%6d %14.4f %14.4f %14.4f'
              % (epoch, tr_loss, tr_right / len(small),
                 va_right / len(val_set)))

parameters 669706 for 300 training images

 epoch     train loss      train acc        val acc
       2.2242         0.1833         0.1615
       0.7274         0.8000         0.6625
       0.1882         0.9633         0.7735
       0.0565         0.9967         0.8020
       0.0263         1.0000         0.8095
       0.0162         1.0000         0.8090

Training accuracy reaches 1.0 and the training loss goes to almost nothing, while validation accuracy stops improving long before that. The model has not learned what a digit is, it has learned these three hundred digits. Week 4 is entirely about the tools for preventing this.

Early stopping, and keeping the right weights

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf)

# Subsets, so every page finishes while you watch it. The lessons are the
# same at full size and the numbers are a little higher.
train_set = Subset(train_all, range(8000))
val_set = Subset(train_all, range(8000, 10000))
test_set = Subset(test_all, range(2000))
def fit(model, epochs=4, lr=0.1, train=None, val=None, log=True):
    """The five steps from week 1, wrapped up so the pages stay short."""
    train = train if train is not None else DataLoader(train_set,
                                                       batch_size=64,
                                                       shuffle=True)
    val = val if val is not None else DataLoader(val_set, batch_size=256)
    loss_fn = nn.CrossEntropyLoss()
    opt = torch.optim.SGD(model.parameters(), lr=lr)
    history = []
    for epoch in range(1, epochs + 1):
        model.train()
        seen = total = 0
        for xb, yb in train:
            opt.zero_grad()
            loss = loss_fn(model(xb), yb)
            loss.backward()
            opt.step()
            total += loss.item() * yb.numel()
            seen += yb.numel()
        model.eval()
        right = count = 0
        with torch.no_grad():
            for xb, yb in val:
                right += (model(xb).argmax(1) == yb).sum().item()
                count += yb.numel()
        history.append((total / seen, right / count))
        if log:
            print('epoch %d  train loss %.4f  val accuracy %.4f'
                  % (epoch, total / seen, right / count))
    return history
import copy
torch.manual_seed(0)

small = Subset(train_all, range(300))
loader = DataLoader(small, batch_size=32, shuffle=True)
val_loader = DataLoader(val_set, batch_size=256)
model = nn.Sequential(nn.Flatten(), nn.Linear(784, 512), nn.ReLU(),
                      nn.Linear(512, 10))
loss_fn = nn.CrossEntropyLoss()
opt = torch.optim.SGD(model.parameters(), lr=0.1)

best_loss, best_epoch, best_state, patience = float('inf'), 0, None, 8
for epoch in range(1, 61):
    model.train()
    for xb, yb in loader:
        opt.zero_grad()
        loss_fn(model(xb), yb).backward()
        opt.step()
    model.eval()
    with torch.no_grad():
        vloss = sum(loss_fn(model(xb), yb).item() * yb.numel()
                    for xb, yb in val_loader) / len(val_set)
    if vloss < best_loss:
        best_loss, best_epoch = vloss, epoch
        best_state = copy.deepcopy(model.state_dict())
    elif epoch - best_epoch >= patience:
        print('stopped at epoch %d' % epoch)
        break

print('best validation loss %.4f at epoch %d' % (best_loss, best_epoch))
model.load_state_dict(best_state)
with torch.no_grad():
    right = sum((model(xb).argmax(1) == yb).sum().item()
                for xb, yb in val_loader)
print('restored model accuracy %.4f' % (right / len(val_set)))

stopped at epoch 27
best validation loss 0.6653 at epoch 19
restored model accuracy 0.7945

Stopping is not the same as restoring

Patience means you keep training for several epochs after the best one, so the weights you are holding when the loop breaks are worse than the ones you had. Saving state_dict() at the best epoch and loading it back is what makes early stopping actually work. copy.deepcopy matters as well: state_dict() returns references to the live tensors, so without the copy your saved best is overwritten by the next optimiser step.

Day 3 takeaway

Three splits: train for gradients, validation for decisions, test for one number at the end. Expect a deep model to reach perfect training accuracy on a small set; that is memorisation, not learning. Deep copy the best state dictionary, because early stopping without restoring is just stopping.

Week 02 · Day 4 of 7

Saving and Reproducing

State dictionaries, checkpoints, and what a seed does not fix

By Pawan Rama Mali766 words

Saving a model sounds trivial and has three separate traps in it: what you save, what you need alongside it, and what happens when the code moves on.

Save the state dictionary, not the model

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf)

# Subsets, so every page finishes while you watch it. The lessons are the
# same at full size and the numbers are a little higher.
train_set = Subset(train_all, range(8000))
val_set = Subset(train_all, range(8000, 10000))
test_set = Subset(test_all, range(2000))
torch.manual_seed(0)
model = nn.Sequential(nn.Flatten(), nn.Linear(784, 32), nn.ReLU(),
                      nn.Linear(32, 10))

state = model.state_dict()
print('state_dict entries:')
for key, value in state.items():
    print('  %-12s %s' % (key, tuple(value.shape)))

torch.save(state, 'model.pt')

fresh = nn.Sequential(nn.Flatten(), nn.Linear(784, 32), nn.ReLU(),
                      nn.Linear(32, 10))
fresh.load_state_dict(torch.load('model.pt', weights_only=True))

x = torch.randn(3, 1, 28, 28)
print('\nsame predictions after reload:',
      bool(torch.allclose(model(x), fresh(x))))

state_dict entries:
  1.weight     (32, 784)
  1.bias       (32,)
  3.weight     (10, 32)
  3.bias       (10,)

same predictions after reload: True

Why not torch.save(model)

Saving the whole object pickles a reference to your class, so the file only loads if that class is importable from the same module path. Rename the file, move the class, or open it in a different project and it breaks. A state dictionary is just named tensors, so it loads anywhere you can construct the same architecture. Pass weights_only=True when loading: without it, torch.load will execute arbitrary code from the file, which matters the moment you download somebody else's checkpoint.

A checkpoint has more in it than weights

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf)

# Subsets, so every page finishes while you watch it. The lessons are the
# same at full size and the numbers are a little higher.
train_set = Subset(train_all, range(8000))
val_set = Subset(train_all, range(8000, 10000))
test_set = Subset(test_all, range(2000))
torch.manual_seed(0)
model = nn.Sequential(nn.Flatten(), nn.Linear(784, 32), nn.ReLU(),
                      nn.Linear(32, 10))
opt = torch.optim.Adam(model.parameters(), lr=1e-3)

# one step, so the optimiser has state worth keeping
loss = nn.CrossEntropyLoss()(model(torch.randn(4, 1, 28, 28)),
                             torch.tensor([1, 2, 3, 4]))
loss.backward()
opt.step()

checkpoint = {
    'epoch': 7,
    'model': model.state_dict(),
    'optimizer': opt.state_dict(),
    'best_val_loss': 0.2431,
    'torch_version': torch.__version__,
}
torch.save(checkpoint, 'checkpoint.pt')

loaded = torch.load('checkpoint.pt', weights_only=False)
print('keys:', sorted(loaded.keys()))
print('resuming from epoch', loaded['epoch'])
print('optimiser state groups:', len(loaded['optimizer']['param_groups']))
print('Adam keeps running averages per parameter, and losing them')
print('makes the first few steps after a resume much worse.')

keys: ['best_val_loss', 'epoch', 'model', 'optimizer', 'torch_version']
resuming from epoch 7
optimiser state groups: 1
Adam keeps running averages per parameter, and losing them
makes the first few steps after a resume much worse.

Reproducibility, honestly

import torch

def two_runs(seeded):
    outs = []
    for _ in range(2):
        if seeded:
            torch.manual_seed(1234)
        model = torch.nn.Linear(4, 2)
        outs.append(model.weight.detach().clone())
    return torch.allclose(outs[0], outs[1])

print('two models, no seed  : identical?', two_runs(False))
print('two models, seeded   : identical?', two_runs(True))

two models, no seed  : identical? False
two models, seeded   : identical? True

Source of variation	Controlled by	Fully solved?
Weight initialisation	`torch.manual_seed`	Yes
Shuffling and augmentation	the same seed, plus a `DataLoader` generator	Yes
Dropout masks	`torch.manual_seed`	Yes
Thread count and reduction order	`torch.set_num_threads`	Mostly
GPU kernel selection	`torch.use_deterministic_algorithms`	At a real cost in speed
A different library version	pinning versions	Only by pinning

Seeded is not the same as reproducible on another machine

Floating point addition is not associative, so a machine that splits a sum across sixteen threads gets a slightly different answer from one that splits it across four. Over thousands of steps those differences compound into visibly different final numbers. Seeds make a run repeatable on your machine, which is what you need for debugging. Matching somebody else's published figure exactly usually needs their hardware as well, and this is why every page on this course pins the thread count and still tells you to expect close rather than identical.

Day 4 takeaway

Save state_dict(), load with weights_only=True, and put the optimiser state, the epoch and the version in the checkpoint alongside the weights. Seed everything, and expect close rather than identical numbers on different hardware.

Week 02 · Day 5 of 7

Habits That Save Days

Parameter budgets, the single batch test, and logging a run

By Pawan Rama Mali1134 words

Some habits that turn a training script from something you babysit into something you can leave running and trust afterwards.

Count your parameters before you train

import torch
from torch import nn

def summarise(model, input_shape):
    total = sum(p.numel() for p in model.parameters())
    trainable = sum(p.numel() for p in model.parameters()
                    if p.requires_grad)
    x = torch.zeros(1, *input_shape)
    print('%-26s %12s %10s' % ('layer', 'output', 'params'))
    for name, mod in model.named_children():
        x = mod(x)
        n = sum(p.numel() for p in mod.parameters())
        print('%-26s %12s %10d'
              % ('%s (%s)' % (name, type(mod).__name__),
                 str(tuple(x.shape[1:])), n))
    print('%-26s %12s %10d' % ('TOTAL', '', total))
    print('trainable %d, frozen %d' % (trainable, total - trainable))
    print('at 4 bytes each that is %.1f MB of weights' % (total * 4 / 1e6))

model = nn.Sequential(nn.Flatten(), nn.Linear(784, 256), nn.ReLU(),
                      nn.Linear(256, 128), nn.ReLU(), nn.Linear(128, 10))
summarise(model, (1, 28, 28))

layer                            output     params
0 (Flatten)                      (784,)          0
1 (Linear)                       (256,)     200960
2 (ReLU)                         (256,)          0
3 (Linear)                       (128,)      32896
4 (ReLU)                         (128,)          0
5 (Linear)                        (10,)       1290
TOTAL                                       235146
trainable 235146, frozen 0
at 4 bytes each that is 0.9 MB of weights

Two useful sanity checks come out of that number. If your parameter count is far larger than your example count, expect to fight overfitting all week. And memory during training is roughly four times the weights, because the gradients, and the optimiser's two running averages if you use Adam, are all the same size as the weights themselves.

Overfit one batch before you trust anything

The single batch test: Take one batch, turn off shuffling and any regularisation, and train on that batch alone. A correctly wired model will drive its loss to nearly zero within a few hundred steps. If it cannot, the bug is in your code, not in your hyperparameters.

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf)

# Subsets, so every page finishes while you watch it. The lessons are the
# same at full size and the numbers are a little higher.
train_set = Subset(train_all, range(8000))
val_set = Subset(train_all, range(8000, 10000))
test_set = Subset(test_all, range(2000))
torch.manual_seed(0)
xb, yb = next(iter(DataLoader(train_set, batch_size=16, shuffle=True)))

model = nn.Sequential(nn.Flatten(), nn.Linear(784, 64), nn.ReLU(),
                      nn.Linear(64, 10))
opt = torch.optim.Adam(model.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()

for step in range(201):
    opt.zero_grad()
    loss = loss_fn(model(xb), yb)
    loss.backward()
    opt.step()
    if step % 50 == 0:
        print('step %3d  loss %.6f  accuracy %.3f'
              % (step, loss.item(),
                 (model(xb).argmax(1) == yb).float().mean().item()))

step   0  loss 2.280022  accuracy 0.375
step  50  loss 0.056201  accuracy 1.000
step 100  loss 0.009544  accuracy 1.000
step 150  loss 0.004860  accuracy 1.000
step 200  loss 0.003012  accuracy 1.000

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf)

# Subsets, so every page finishes while you watch it. The lessons are the
# same at full size and the numbers are a little higher.
train_set = Subset(train_all, range(8000))
val_set = Subset(train_all, range(8000, 10000))
test_set = Subset(test_all, range(2000))
torch.manual_seed(0)
xb, yb = next(iter(DataLoader(train_set, batch_size=16, shuffle=True)))

# The same test on a model with a bug: the labels are shuffled, so there
# is no relationship left to learn.
shuffled = yb[torch.randperm(len(yb))]
model = nn.Sequential(nn.Flatten(), nn.Linear(784, 64), nn.ReLU(),
                      nn.Linear(64, 10))
opt = torch.optim.Adam(model.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()
for step in range(201):
    opt.zero_grad()
    loss_fn(model(xb), shuffled).backward()
    opt.step()
print('with randomised labels, final loss %.6f'
      % loss_fn(model(xb), shuffled).item())
print('it still memorises 16 examples, which is the point:')
print('the test proves your wiring works, not that your data means anything.')

with randomised labels, final loss 0.004961
it still memorises 16 examples, which is the point:
the test proves your wiring works, not that your data means anything.

Log what you will want later

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf)

# Subsets, so every page finishes while you watch it. The lessons are the
# same at full size and the numbers are a little higher.
train_set = Subset(train_all, range(8000))
val_set = Subset(train_all, range(8000, 10000))
test_set = Subset(test_all, range(2000))
def fit(model, epochs=4, lr=0.1, train=None, val=None, log=True):
    """The five steps from week 1, wrapped up so the pages stay short."""
    train = train if train is not None else DataLoader(train_set,
                                                       batch_size=64,
                                                       shuffle=True)
    val = val if val is not None else DataLoader(val_set, batch_size=256)
    loss_fn = nn.CrossEntropyLoss()
    opt = torch.optim.SGD(model.parameters(), lr=lr)
    history = []
    for epoch in range(1, epochs + 1):
        model.train()
        seen = total = 0
        for xb, yb in train:
            opt.zero_grad()
            loss = loss_fn(model(xb), yb)
            loss.backward()
            opt.step()
            total += loss.item() * yb.numel()
            seen += yb.numel()
        model.eval()
        right = count = 0
        with torch.no_grad():
            for xb, yb in val:
                right += (model(xb).argmax(1) == yb).sum().item()
                count += yb.numel()
        history.append((total / seen, right / count))
        if log:
            print('epoch %d  train loss %.4f  val accuracy %.4f'
                  % (epoch, total / seen, right / count))
    return history
import json
torch.manual_seed(0)

model = nn.Sequential(nn.Flatten(), nn.Linear(784, 64), nn.ReLU(),
                      nn.Linear(64, 10))
history = fit(model, epochs=3, log=False)

run = {
    'seed': 0,
    'architecture': '784-64-10',
    'optimizer': 'SGD',
    'lr': 0.1,
    'batch_size': 64,
    'epochs': len(history),
    'train_examples': len(train_set),
    'history': [{'epoch': i + 1, 'train_loss': round(l, 4),
                 'val_acc': round(a, 4)}
                for i, (l, a) in enumerate(history)],
}
print(json.dumps(run, indent=2)[:600])

{
  "seed": 0,
  "architecture": "784-64-10",
  "optimizer": "SGD",
  "lr": 0.1,
  "batch_size": 64,
  "epochs": 3,
  "train_examples": 8000,
  "history": [
    {
      "epoch": 1,
      "train_loss": 1.0652,
      "val_acc": 0.862
    },
    {
      "epoch": 2,
      "train_loss": 0.416,
      "val_acc": 0.8885
    },
    {
      "epoch": 3,
      "train_loss": 0.3351,
      "val_acc": 0.897
    }
... (2 more lines)

The number you did not record is the run you cannot repeat

Six weeks into a project you will have thirty runs and a vague memory that one of them worked. Write the configuration and the per-epoch numbers to a file named after the run, every time, automatically. Tools such as TensorBoard, Weights and Biases or MLflow do this and draw the curves for you, and all of them are better than the notebook cell you will overwrite this afternoon.

Day 5 takeaway

Count parameters before training and budget about four times that in memory. Overfit a single batch to prove the wiring works before you blame a hyperparameter. Log the configuration with the results, in a file, automatically.

Week 02 · Day 6 of 7

Real Data Problems

Normalisation, and batches whose examples are different lengths

By Pawan Rama Mali800 words

Two things that will bite you on real data rather than on MNIST: input scaling, and batches whose contents are not all the same size.

Normalising the input

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf)

# Subsets, so every page finishes while you watch it. The lessons are the
# same at full size and the numbers are a little higher.
train_set = Subset(train_all, range(8000))
val_set = Subset(train_all, range(8000, 10000))
test_set = Subset(test_all, range(2000))
raw = train_all.data[:8000].float() / 255.0
print('pixel mean %.4f, std %.4f' % (raw.mean(), raw.std()))

normalised = (raw - raw.mean()) / raw.std()
print('after normalising: mean %.4f, std %.4f'
      % (normalised.mean(), normalised.std()))

print('\ntorchvision does this for you:')
print("  transforms.Normalize((0.1307,), (0.3081,))")
print('  which are the mean and standard deviation of MNIST.')

pixel mean 0.1308, std 0.3082
after normalising: mean -0.0000, std 1.0000

torchvision does this for you:
  transforms.Normalize((0.1307,), (0.3081,))
  which are the mean and standard deviation of MNIST.

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf)

# Subsets, so every page finishes while you watch it. The lessons are the
# same at full size and the numbers are a little higher.
train_set = Subset(train_all, range(8000))
val_set = Subset(train_all, range(8000, 10000))
test_set = Subset(test_all, range(2000))
torch.manual_seed(0)

def train_briefly(transform):
    data = datasets.MNIST('data', train=True, download=True,
                          transform=transform)
    loader = DataLoader(Subset(data, range(4000)), batch_size=64,
                        shuffle=True)
    torch.manual_seed(0)
    model = nn.Sequential(nn.Flatten(), nn.Linear(784, 64), nn.ReLU(),
                          nn.Linear(64, 10))
    opt = torch.optim.SGD(model.parameters(), lr=0.05)
    loss_fn = nn.CrossEntropyLoss()
    losses = []
    for epoch in range(3):
        for xb, yb in loader:
            opt.zero_grad()
            loss = loss_fn(model(xb), yb)
            loss.backward()
            opt.step()
        losses.append(loss.item())
    return losses

plain = train_briefly(transforms.ToTensor())
normed = train_briefly(transforms.Compose([
    transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]))
print('%-14s %s' % ('0 to 1', ['%.4f' % v for v in plain]))
print('%-14s %s' % ('normalised', ['%.4f' % v for v in normed]))

0 to 1         ['1.4501', '0.7852', '0.7656']
normalised     ['0.5104', '0.4380', '0.5037']

Compute the statistics on training data only

The mean and standard deviation are learned quantities, and computing them over the whole dataset lets the test set influence how the training set is scaled. It is a small leak and it is the same leak as any other. Use the training statistics for every split, including at serving time, where they have to be shipped alongside the weights.

Batches of unequal length

import torch
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import DataLoader

sequences = [torch.tensor([1, 2, 3]),
             torch.tensor([4, 5]),
             torch.tensor([6, 7, 8, 9, 10])]

try:
    torch.stack(sequences)
except RuntimeError as e:
    print('stacking them directly fails:')
    print(' ', str(e)[:78])

padded = pad_sequence(sequences, batch_first=True, padding_value=0)
print('\npadded to a rectangle:')
print(padded)
print('lengths', [len(s) for s in sequences])
mask = padded != 0
print('mask (True means a real token):')
print(mask)

stacking them directly fails:
  stack expects each tensor to be equal size, but got [3] at entry 0 and [2] at 

padded to a rectangle:
tensor([[ 1,  2,  3,  0,  0],
        [ 4,  5,  0,  0,  0],
        [ 6,  7,  8,  9, 10]])
lengths [3, 2, 5]
mask (True means a real token):
tensor([[ True,  True,  True, False, False],
        [ True,  True, False, False, False],
        [ True,  True,  True,  True,  True]])

import torch
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import DataLoader, Dataset

class Variable(Dataset):
    def __len__(self):
        return 6

    def __getitem__(self, i):
        return torch.arange(1, i + 2), i % 2

def collate(batch):
    """Called with a list of examples, returns one batch."""
    seqs, labels = zip(*batch)
    lengths = torch.tensor([len(s) for s in seqs])
    return (pad_sequence(seqs, batch_first=True),
            torch.tensor(labels), lengths)

loader = DataLoader(Variable(), batch_size=3, collate_fn=collate)
for x, y, lengths in loader:
    print('batch shape', tuple(x.shape), ' lengths', lengths.tolist())
    print(x)

batch shape (3, 3)  lengths [1, 2, 3]
tensor([[1, 0, 0],
        [1, 2, 0],
        [1, 2, 3]])
batch shape (3, 6)  lengths [4, 5, 6]
tensor([[1, 2, 3, 4, 0, 0],
        [1, 2, 3, 4, 5, 0],
        [1, 2, 3, 4, 5, 6]])

collate_fn is the hook for anything that cannot be stacked automatically: variable-length text, images of different sizes, examples that need grouping. Weeks 8 and 9 rely on it heavily, and the padding mask it produces is the same mask that attention will need.

Day 6 takeaway

Normalise inputs using training-set statistics, and ship those statistics with the model. When examples differ in length, pad them in a collate_fn and carry a mask alongside, because every layer downstream needs to know which positions are real.

Week 02 · Day 7 of 7

A Complete Training Script

Everything assembled, with the deliberate omissions named

By Pawan Rama Mali827 words

Everything from this week in one script: a proper module, a real validation loop, early stopping with restoration, checkpointing and a run record.

The model

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf)

# Subsets, so every page finishes while you watch it. The lessons are the
# same at full size and the numbers are a little higher.
train_set = Subset(train_all, range(8000))
val_set = Subset(train_all, range(8000, 10000))
test_set = Subset(test_all, range(2000))
class MLP(nn.Module):
    def __init__(self, hidden=(256, 128), n_classes=10):
        super().__init__()
        sizes = (784,) + tuple(hidden)
        self.blocks = nn.ModuleList([
            nn.Sequential(nn.Linear(a, b), nn.ReLU())
            for a, b in zip(sizes, sizes[1:])])
        self.head = nn.Linear(sizes[-1], n_classes)

    def forward(self, x):
        x = x.flatten(1)
        for block in self.blocks:
            x = block(x)
        return self.head(x)

model = MLP()
print(model)
print('\nparameters %d' % sum(p.numel() for p in model.parameters()))

MLP(
  (blocks): ModuleList(
    (0): Sequential(
      (0): Linear(in_features=784, out_features=256, bias=True)
      (1): ReLU()
    )
    (1): Sequential(
      (0): Linear(in_features=256, out_features=128, bias=True)
      (1): ReLU()
    )
  )
  (head): Linear(in_features=128, out_features=10, bias=True)
)

parameters 235146

The training script

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf)

# Subsets, so every page finishes while you watch it. The lessons are the
# same at full size and the numbers are a little higher.
train_set = Subset(train_all, range(8000))
val_set = Subset(train_all, range(8000, 10000))
test_set = Subset(test_all, range(2000))
import copy, json, time

class MLP(nn.Module):
    def __init__(self, hidden=(256, 128)):
        super().__init__()
        sizes = (784,) + tuple(hidden)
        self.blocks = nn.ModuleList([
            nn.Sequential(nn.Linear(a, b), nn.ReLU())
            for a, b in zip(sizes, sizes[1:])])
        self.head = nn.Linear(sizes[-1], 10)

    def forward(self, x):
        x = x.flatten(1)
        for block in self.blocks:
            x = block(x)
        return self.head(x)

CONFIG = {'seed': 0, 'hidden': (256, 128), 'lr': 0.1, 'batch_size': 64,
          'max_epochs': 30, 'patience': 5}

torch.manual_seed(CONFIG['seed'])
train_loader = DataLoader(train_set, batch_size=CONFIG['batch_size'],
                          shuffle=True)
val_loader = DataLoader(val_set, batch_size=256)
test_loader = DataLoader(test_set, batch_size=256)

model = MLP(CONFIG['hidden'])
opt = torch.optim.SGD(model.parameters(), lr=CONFIG['lr'])
loss_fn = nn.CrossEntropyLoss()

def evaluate(loader):
    model.eval()
    loss = right = seen = 0
    with torch.no_grad():
        for xb, yb in loader:
            out = model(xb)
            loss += loss_fn(out, yb).item() * yb.numel()
            right += (out.argmax(1) == yb).sum().item()
            seen += yb.numel()
    return loss / seen, right / seen

best = {'loss': float('inf'), 'epoch': 0, 'state': None}
history = []
start = time.time()
for epoch in range(1, CONFIG['max_epochs'] + 1):
    model.train()
    for xb, yb in train_loader:
        opt.zero_grad()
        loss_fn(model(xb), yb).backward()
        opt.step()
    vloss, vacc = evaluate(val_loader)
    history.append({'epoch': epoch, 'val_loss': round(vloss, 4),
                    'val_acc': round(vacc, 4)})
    if vloss < best['loss']:
        best = {'loss': vloss, 'epoch': epoch,
                'state': copy.deepcopy(model.state_dict())}
    elif epoch - best['epoch'] >= CONFIG['patience']:
        break

model.load_state_dict(best['state'])
test_loss, test_acc = evaluate(test_loader)

torch.save({'config': CONFIG, 'model': best['state'],
            'best_epoch': best['epoch'], 'test_acc': test_acc},
           'mnist_mlp.pt')

print('ran %d epochs in %.1fs' % (len(history), time.time() - start))
print('best epoch %d, validation loss %.4f'
      % (best['epoch'], best['loss']))
print('test accuracy %.4f  (looked at once)' % test_acc)
print('\nlast five epochs:')
for row in history[-5:]:
    print(' ', json.dumps(row))

ran 22 epochs in 20.4s
best epoch 17, validation loss 0.2238
test accuracy 0.9225  (looked at once)

last five epochs:
  {"epoch": 18, "val_loss": 0.2358, "val_acc": 0.936}
  {"epoch": 19, "val_loss": 0.2267, "val_acc": 0.939}
  {"epoch": 20, "val_loss": 0.2451, "val_acc": 0.934}
  {"epoch": 21, "val_loss": 0.2329, "val_acc": 0.939}
  {"epoch": 22, "val_loss": 0.2314, "val_acc": 0.939}

The habits, collected

Subclass nn.Module; use nn.ModuleList for anything you build in a loop.
Pass logits to the loss and let it apply the softmax.
Three splits, and look at the test set once.
Overfit a single batch before you debug anything else.
Deep copy the best state dictionary and restore it at the end.
Save the configuration in the checkpoint next to the weights.
Normalise inputs with training statistics and ship them with the model.
Log per-epoch numbers to a file, not to a scrollback buffer.

What is missing, deliberately

This model has no dropout, no weight decay, no augmentation, no learning rate schedule and no normalisation layers. It is a clean baseline, and it is what the next two weeks improve on one change at a time so you can see what each is worth. A model with all five added at once is a model where you cannot tell which of them helped, and in practice one or two usually do nothing at all.

Day 7 takeaway

You can now build an arbitrary architecture, choose a loss that matches the task, split data honestly, stop at the right epoch and restore the right weights, and save something you can come back to in six months. Week 3 makes it train faster and week 4 makes it generalise better.