🎓 A1Learner
🔍

PyTorch Model Capacity

Model capacity is how much pattern a model can hold, set mostly by its size. Too little capacity underfits, and too much capacity can memorise the training rows.

Width is parameter count

Width sets the size of the hidden layers. The parameter count grows fast as width grows.

import torch

for width in [2, 16, 256]:
    model = torch.nn.Sequential(
        torch.nn.Linear(1, width), torch.nn.ReLU(),
        torch.nn.Linear(width, width), torch.nn.ReLU(),
        torch.nn.Linear(width, 1))
    count = sum(p.numel() for p in model.parameters())
    print(f"Width {width}: {count} parameters")

One problem, three sizes

A fair size test changes only the width. Data, steps, and the seed stay fixed across the runs.

for width in [2, 16, 256]:
    model = train(width)  # same data, same steps, only width changes

Reading the result

Training loss and validation loss tell different stories. Both are needed before picking a size.

train loss falls with width  ->  the big model fits the rows
val loss rises again at 256  ->  the fit did not travel to new rows

Example

import torch

torch.manual_seed(0)
x = torch.rand(160, 1) * 4 - 2
y = torch.sin(2 * x) + 0.3 * x + torch.randn(160, 1) * 0.4
x_train, x_val = x[:24], x[24:]
y_train, y_val = y[:24], y[24:]
loss_fn = torch.nn.MSELoss()


def train(width):
    torch.manual_seed(1)
    model = torch.nn.Sequential(
        torch.nn.Linear(1, width), torch.nn.ReLU(),
        torch.nn.Linear(width, width), torch.nn.ReLU(),
        torch.nn.Linear(width, 1))
    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
    for step in range(400):
        optimizer.zero_grad()
        loss = loss_fn(model(x_train), y_train)
        loss.backward()
        optimizer.step()
    with torch.no_grad():
        train_loss = loss_fn(model(x_train), y_train).item()
        val_loss = loss_fn(model(x_val), y_val).item()
    return train_loss, val_loss


for width in [2, 16, 256]:
    train_loss, val_loss = train(width)
    print(f"Width {width}: train loss {train_loss:.4f}, "
          f"val loss {val_loss:.4f}")

Key points

  • Capacity: how much pattern a model can hold, growing with its parameter count
  • Underfit: to miss the real pattern, so both training and validation losses stay high
  • Memorise: to fit the training rows themselves, noise included, instead of the pattern
  • Validation data: rows held back from training, used to check the model on unseen examples
  • Use the validation loss to compare sizes, because training loss tends to favour the bigger model
  • Capacity is a knob to set per problem, and the middle setting often generalises best

In this example, 24 rows train each model and 136 rows wait as validation data. The targets follow a sine curve with noise added. Three widths train for the same 400 steps from the same seed. Width 2 underfits, with train loss 0.0894 and validation loss 0.2019. Width 16 fits better and generalises best, with validation loss 0.1700. Width 256 reaches the lowest train loss, 0.0256, and the worst validation loss, 0.2351. The widest model fit the training rows themselves, noise included.

Note

A bigger model is not a better model by default, and the validation loss is the judge.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

A model has the lowest training loss and the highest validation loss in a size test. What does that suggest?