Network design

Overview

The Bouncing ball lesson made five choices before training even started:

Training never touches these numbers; we set them by hand. Choices like these are called hyperparameters, to tell them apart from the weights and biases that training learns.

Every test follows the same three rules:

• Change one thing.
• Try three random starts. The network’s first weights are random, like a shuffled deck, so we shuffle three times to see how much of a result is luck.
• Grade on the 100 test throws the network never saw.

The grades are two misses, both for a single 10 ms guess:

• the miss in the air: how far off the guessed position is while the ball flies, in millimetres;
• the miss at the bounce: how far off the guessed speed is at the floor, in metres per second.

Network size

Each neuron adds one bend the network can use. So how many bends does a bounce need? Here is how we trained networks of different sizes, each from three random starts:

python3 03_train.py --hidden 2 --seed 0     # 2 neurons, random start 1
python3 03_train.py --hidden 2 --seed 1     # 2 neurons, random start 2
python3 03_train.py --hidden 16 --seed 2    # 16 neurons, random start 3
python3 04_evaluate.py --tag hidden16_lr0.1_seed2    # grade it on the test throws

--hidden sets how many neurons make_network puts in the hidden layer, and --seed picks the shuffle for the starting weights:

torch.manual_seed(args.seed)               # which random starting weights
network = make_network(args.hidden[0] if len(args.hidden) == 1 else args.hidden)
# --hidden 2 gives nn.Linear(4, 2), nn.ReLU(), nn.Linear(2, 4)

"""
Step 3 — Train the network to predict the change in 10 ms.

    input  : s_t = [x, z, vx, vz]            output : s_(t+1) - s_t
    model  : 4 -> 16 ReLU neurons -> 4        loss   : mean squared error, in raw units (m, m/s)
    update : plain SGD, learning rate 0.1, batches of 256, 240 passes over the data

Run:  python 03_train.py                      (the lesson's network)
      python 03_train.py --hidden 2           (change one thing; see 05_compare.py)
      python 03_train.py --hidden 8 8         (two hidden layers of 8)
Out:  results/<tag>/network.pt, results/<tag>/losses.json
"""

import argparse, json, os
import numpy as np
import torch
from torch import nn
from model import make_network

parser = argparse.ArgumentParser()
parser.add_argument("--hidden", type=int, nargs="+", default=[16])
parser.add_argument("--lr", type=float, default=0.1)
parser.add_argument("--seed", type=int, default=0)
parser.add_argument("--epochs", type=int, default=240)
parser.add_argument("--batch", type=int, default=256)
args = parser.parse_args()
EPOCHS, BATCH = args.epochs, args.batch
tag = f"hidden{'x'.join(map(str, args.hidden))}_lr{args.lr:g}_seed{args.seed}"
tag += (f"_e{EPOCHS}" if EPOCHS != 240 else "") + (f"_b{BATCH}" if BATCH != 256 else "")

# Data: every pair from the training throws.
data = np.load("data/throws.npz")
throws = data["states"][data["split"] == "train"]             # (500, 201, 4)
now, after = throws[:, :-1].reshape(-1, 4), throws[:, 1:].reshape(-1, 4)
inputs = torch.tensor(now, dtype=torch.float32)               # (100000, 4)
targets = torch.tensor(after - now, dtype=torch.float32)      # the change in 10 ms

torch.manual_seed(args.seed)
network = make_network(args.hidden[0] if len(args.hidden) == 1 else args.hidden)
loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(network.parameters(), lr=args.lr)

losses = []
for epoch in range(EPOCHS):
    order = torch.randperm(len(inputs))                       # a new shuffle every pass
    for i in range(0, len(order), BATCH):
        batch = order[i:i + BATCH]
        loss = loss_fn(network(inputs[batch]), targets[batch])
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    with torch.no_grad():
        losses.append(loss_fn(network(inputs), targets).item())
    if (epoch + 1) % 40 == 0:
        print(f"epoch {epoch + 1:3d}  loss {losses[-1]:.5f}")

os.makedirs(f"results/{tag}", exist_ok=True)
torch.save({"hidden": args.hidden[0] if len(args.hidden) == 1 else args.hidden, "weights": network.state_dict()}, f"results/{tag}/network.pt")
json.dump({"args": vars(args), "losses": losses}, open(f"results/{tag}/losses.json", "w"))
print(f"saved results/{tag}/")

05_compare.py runs all twelve (four sizes × three starts) for you. First, how the mistake shrank during training, one line per random start:

A line that stays high never learned; a line that drops did.

• With 2 and 4 neurons, it depends on luck: some starts get stuck near the top, others drop.
• With 8, every start learns something.
• With 16, all three lines land in the same low place: luck no longer decides.

But a low loss can fool you. One of the 2-neuron networks ends almost as low as 16 neurons, yet it cannot fly.

Why? The bounce’s mistakes are huge and the flying mistakes tiny, so the loss is mostly about the bounce. That is why we grade the two jobs separately:

The dashed line is plain physics with no floor, which misses every bounce completely.

Small networks are all over the place; 16 neurons sit together near the bottom.

To see why the small ones fail, look inside them at their two jobs:

• in the air, the height change should follow the speed in a straight line;
• near the floor, the speed should suddenly jump when the ball is low and falling.

The dashed line is the right answer.

• With 2 neurons there are only two bends to spend, and training spends them on the bounce, where the mistakes are biggest. Flying gets nothing: from starts 1 and 3, the height change ignores the speed completely, a ball that cannot fly.
• With 4, two starts get stuck and one lucky start learns both jobs.
• With 8, every start learns the bounce, but flying only sometimes.
• With 16, every start learns both.

Finally, the same test throw, played out by each network on its own for two seconds:

Blue: each network predicting the throw on its own from the same start (random start 1).

So: 16 is the smallest network that learned both jobs from every start. Bigger would cost more computing for the same result.

A good rule of thumb in general: pick the smallest network that works every time you try it.

Learning rate

The learning rate is the size of each step. Too small and training crawls; too big and it jumps past the answer, like the loss bowl in PyTorch basics. To try another one, change one number:

python3 03_train.py --lr 0.01    # ten times smaller steps
python3 03_train.py --lr 0.3     # three times bigger steps

optimizer = torch.optim.SGD(network.parameters(), lr=args.lr)

Same 16-neuron network, same data, three step sizes:

Only the step size changed.

• At 0.01 the steps are so small that after 240 epochs the network has barely found the bounce.
• At 0.3 they are so big that it keeps jumping past the answer and never settles. Played out on its own, its ball flies off by kilometres.
• 0.1 learns.

The spikes in its curve are batches that happened to hold several bounces: big mistakes, so a big step, and the next batches pull it back.

Spikes like that mean 0.1 is close to too big.

So: 0.1 is the biggest step that still settles.

Batch size

Each step averages the corrections of a batch of examples (Gradients and backpropagation explains why). Is a bigger batch better?

Here is the fair race: four runs, each given exactly 20 seconds of computer time, with batches of 1, 32, 256, or all 100,000 pairs at once:

python3 06_batch_size.py 32      # one racer: batches of 32 for 20 seconds
python3 06_batch_size.py         # all four racers at once, then grade them

for i in range(0, len(order), batch):     # walk through the shuffled pairs, batch at a time
    rows = order[i:i + batch]             # e.g. 32 rows -> one step

"""
Step 6 — Batch size, on an equal time budget.

Each run trains the lesson's network (4 -> 16 -> 4, SGD, lr 0.1) for the same 20 seconds of computer time, with
batches of 1, 32, 256, or all 100,000 pairs. It records the loss on all training pairs as it goes, and how many
updates it managed. A fair race: small batches make many cheap, noisy steps; big batches make few, accurate ones.

Run:  python 06_batch_size.py      (about 1.5 minutes; the four runs go in parallel, one CPU core each)
Out:  results/batch<B>_t20/network.pt, losses.json (+ 04_evaluate.py metrics)
"""

import json, os, subprocess, sys, time
import numpy as np
import torch
from torch import nn
from model import make_network

SECONDS, LR = 20.0, 0.1


def race(batch):
    torch.set_num_threads(1)
    torch.manual_seed(0)
    data = np.load("data/throws.npz")
    throws = data["states"][data["split"] == "train"]
    now, after = throws[:, :-1].reshape(-1, 4), throws[:, 1:].reshape(-1, 4)
    inputs = torch.tensor(now, dtype=torch.float32)
    targets = torch.tensor(after - now, dtype=torch.float32)
    network, loss_fn = make_network(16), nn.MSELoss()
    optimizer = torch.optim.SGD(network.parameters(), lr=LR)

    record, updates, trained, next_check = [], 0, 0.0, 0.0     # trained = training time only (checks not counted)
    while trained < SECONDS:
        order = torch.randperm(len(inputs))
        for i in range(0, len(order), batch):
            if trained >= next_check:                           # measure the loss every 0.25 s of training
                with torch.no_grad():
                    record.append([round(trained, 3), updates, loss_fn(network(inputs), targets).item()])
                next_check += 0.25
            start = time.perf_counter()
            rows = order[i:i + batch]
            loss = loss_fn(network(inputs[rows]), targets[rows])
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            trained += time.perf_counter() - start
            updates += 1
            if trained >= SECONDS:
                break
    with torch.no_grad():
        record.append([round(trained, 3), updates, loss_fn(network(inputs), targets).item()])
    tag = f"batch{batch}_t20"
    os.makedirs(f"results/{tag}", exist_ok=True)
    torch.save({"hidden": 16, "weights": network.state_dict()}, f"results/{tag}/network.pt")
    json.dump({"batch": batch, "seconds": SECONDS, "updates": updates, "record": record}, open(f"results/{tag}/losses.json", "w"))
    print(f"batch {batch:6d}: {updates:7d} updates in {SECONDS:.0f} s, final loss {record[-1][2]:.5f}")
    return tag


if __name__ == "__main__":
    if len(sys.argv) > 1:                                       # one race (used by the parallel launcher below)
        race(int(sys.argv[1]))
    else:
        env = {**os.environ, "OMP_NUM_THREADS": "1"}
        procs = [subprocess.Popen(["python3", "06_batch_size.py", str(b)], env=env) for b in (1, 32, 256, 100000)]
        for proc in procs:
            proc.wait()
        for b in (1, 32, 256, 100000):
            subprocess.run(["python3", "04_evaluate.py", "--tag", f"batch{b}_t20"], check=True, env=env)

Two surprises:

• Batches of 32 made almost as many steps as batches of 1. A computer takes about the same time for one step whether it looks at 1 example or 32, so a batch of 32 learns from 32 times more examples in the same time.
• Batches of 1 made plenty of steps but never learned the bounce. Each single bounce example shoves the weights hard its own way, and the next example shoves them back. Averaging a batch calms the shoving.

And the other end? With all 100,000 pairs per step, only 1,664 steps fit into 20 seconds.

That racer actually saw the most examples of all, 166 million, and still learned the least. Learning needs many steps, not just lots of data.

So: batches in the middle, here 32 or 256, give many steps that are each calm enough to trust.

Training time

One epoch is one trip through all 100,000 training pairs. How many trips are enough? Same network, trained for different lengths:

python3 03_train.py --epochs 60      # a quarter of the training
python3 03_train.py --epochs 480     # twice the training

for epoch in range(EPOCHS):                # EPOCHS comes from --epochs (240 if not given)
    order = torch.randperm(len(inputs))    # a new shuffle for every trip

From 30 to 480 epochs:

• the miss in the air shrinks from 6.3 mm to 1.5 mm;
• the miss at the bounce from 2.0 to 0.6 m/s.

But each doubling buys less than the one before. Our 240 epochs, about 20 seconds, is a sensible place to stop, not the best possible result.

Could training too long make things worse? It can, when a network starts memorizing its training examples instead of learning the physics, like a student who memorizes last year’s exam answers and then fails a new exam.

Overfitting

memorizing the training examples instead of learning the rule. You spot it when the test misses get worse while training keeps improving.

Here they keep improving together: with 100,000 examples and only 148 numbers to adjust, there is nothing to memorize.

So: train until the improvement gets too small to be worth the wait.

One or two layers

A second hidden layer reads the first layer’s ramps and bends them. That could help with the bounce, which needs a corner: the speed flips only when the ball is low and falling.

One neuron in the first layer could notice “low”, another “falling”, and a neuron in the second layer could switch on only when both are true.

To keep the test fair, both networks get the same number of weights and biases, 148:

python3 03_train.py --hidden 16      # one hidden layer of 16 neurons
python3 03_train.py --hidden 8 8     # two hidden layers of 8 neurons

for size in sizes:                                # [16] gives one layer; [8, 8] gives two
    layers += [nn.Linear(width, size), nn.ReLU()]
    width = size
return nn.Sequential(*layers, nn.Linear(width, 4))

"""
model.py — the learned part: a network that maps a state to its change over the next 10 ms.
One hidden layer of ReLU neurons gives the rule the bends it needs to turn a fall into a bounce.
"""

import torch
from torch import nn


def make_network(hidden=16):
    """4 numbers in (the state) -> hidden layer(s) of ReLU neurons -> 4 numbers out (the change).

    hidden=16 gives one hidden layer of 16 neurons; hidden=[8, 8] gives two hidden layers of 8.
    """
    sizes = [hidden] if isinstance(hidden, int) else list(hidden)
    layers, width = [], 4
    for size in sizes:
        layers += [nn.Linear(width, size), nn.ReLU()]
        width = size
    return nn.Sequential(*layers, nn.Linear(width, 4))


def load_network(path):
    saved = torch.load(path)
    network = make_network(saved["hidden"])
    network.load_state_dict(saved["weights"])
    return network


def predict_next(network, states):
    """Next state = current state + the network's guessed change. Works on one state or many."""
    with torch.no_grad():
        change = network(torch.tensor(states, dtype=torch.float32)).numpy()
    return states + change

Here is a map of where each network predicts a bounce, for balls at different heights and speeds near the floor. The brighter the colour, the bigger the upward kick:

The bounce lives in the bottom-left corner of each map, where the ball is low and falling fast.

What the second layer bought and cost:

• Better: it draws the corner more sharply. Its miss at the bounce is smaller from every random start.
• Better: for a ball resting on the floor, it predicts a much smaller false kick (0.6 against 1.1 m/s, in the row of zero speed near the top left of each map).
• Worse: it flies much worse, and when it plays out a throw on its own, it drifts far more.

With only 8 neurons per layer and the same training, it could not also learn the straight-line flying as well.

So: for this problem, one wide layer wins overall. Extra layers pay off when a problem is built from pieces of pieces, like edges that form shapes that form a face in a photo.

Recap

The numbers matter less than the method: change one thing, try a few random starts, and grade on examples the network never saw.