Saving and loading models (checkpoints): Exercises

These problems climb from "can you recognise the pieces" up to "can you design a checkpoint system that never loses work." Every one has a full worked solution hidden in a collapsible callout — read the problem, try it yourself, then reveal. This is the parent Saving and loading models (checkpoints) made testable.

Level 1 — Recognition

Exercise 1.1 (L1)

Name the five things that make up a full training state, and mark which single one is not needed if you only ever want to run predictions (inference), never resume training.

Recall Solution 1.1

The five pieces of training state:

Model parameters (weights $W$ , biases $b$ )
Optimizer state (momentum / running averages)
Epoch / step counter
Random seeds
Model architecture (optional blueprint)

For inference only you can drop the optimizer state (and the epoch counter and seeds too, in practice). You keep parameters and architecture — those alone determine predictions. The optimizer state only matters to continue training.

Exercise 1.2 (L1)

In PyTorch, what does model.state_dict() return — the whole nn.Module object, or something smaller? Why does PyTorch prefer to save that smaller thing?

Recall Solution 1.2

It returns a dictionary mapping each layer's name (e.g. "0.weight") to its parameter tensor — not the whole module object.

Why: separation of architecture from parameters. The model code (the class) can evolve, but an old checkpoint still loads as long as the layer names and tensor shapes match. Saving the whole object would tie the file to a specific class definition and pickle version.

Exercise 1.3 (L1)

Match each Keras save option to its consequence: (a) save_weights_only=True → ? (b) model.save('my_model') (SavedModel) → ?

Recall Solution 1.3

(a) save_weights_only=True saves only the weights (binary + index). Smaller, faster, but you must re-define the architecture in code before loading.
(b) model.save('my_model') saves the full model — architecture + weights + computational graph — into a directory. No original code needed; ideal for Model Deployment.

Level 2 — Application

Exercise 2.1 (L2)

A 500-million-parameter model is stored in float32 (4 bytes per parameter). Ignoring optimizer state, roughly how large is one weights-only checkpoint in gigabytes (use $1 GB = 1 0^{9}$ bytes)?

Recall Solution 2.1

$bytes = 500 \times 1 0^{6} \times 4 = 2 \times 1 0^{9} bytes = 2 GB .$ So one checkpoint $\approx 2 GB$ . This is exactly the "~2GB per checkpoint" figure the parent note used.

Exercise 2.2 (L2)

Using the tradeoff equation $Storage = Model Size \times Checkpoint Frequency \times Training Duration,$ compute total storage if you save that 2 GB model every epoch for 100 epochs. Then, if you switch to "keep only the last 3 checkpoints," how much do you store instead?

Recall Solution 2.2

Save-every-epoch: $2 GB \times 1 (per epoch) \times 100 epochs = 200 GB .$ Rolling window of 3: $2 GB \times 3 = 6 GB .$ A 33× reduction in disk. The rolling window trades away distant history you almost never need.

Exercise 2.3 (L2)

You crashed during epoch 7 (epoch 7 never finished, so no epoch-7 checkpoint was written; the last file on disk is checkpoint_0006.pt). Your resume code does start_epoch = checkpoint['epoch'] + 1. What value does start_epoch take, and is that correct?

Recall Solution 2.3

The loaded checkpoint has epoch = 6, so $start_epoch = 6 + 1 = 7.$ Correct. Epoch 6 is the last fully completed epoch, so training resumes at epoch 7 — the one that was interrupted. No completed work is repeated, and no epoch is skipped.

Level 3 — Analysis

Exercise 3.1 (L3)

SGD with momentum uses $v_{t} = β v_{t - 1} + (1 - β) \nabla_{θ} L, θ_{t} = θ_{t - 1} - α v_{t} .$ Suppose every gradient is the same constant $g$ (a long flat slope). Starting from $v_{0} = 0$ , the velocity approaches a steady value $v_{\infty}$ . Find $v_{\infty}$ , and explain what "losing momentum" costs in terms of step size.

Figure — Saving and loading models (checkpoints)

Recall Solution 3.1

At steady state $v_{\infty} = v_{t} = v_{t - 1}$ , so $v_{\infty} = β v_{\infty} + (1 - β) g \Rightarrow v_{\infty} (1 - β) = (1 - β) g \Rightarrow v_{\infty} = g .$ Interesting — the steady velocity equals $g$ itself (this is the " $(1 - β)$ normalised" form). The point is the build-up: starting from $0$ , $v_{t} = g (1 - β^{t})$ climbs toward $g$ over many steps (look at the red curve rising in the figure).

Cost of losing momentum: on resume with $v = 0$ , your first step is $α (1 - β) g$ instead of the cruising $α g$ — only a $(1 - β) = 0.1$ fraction of full step size at $β = 0.9$ . You crawl until the buffer refills.

Exercise 3.2 (L3)

The parent note says momentum takes " $\sim 1/ (1 - β)$ steps to rebuild." For $β = 0.9$ that is $10$ steps (the note rounds to 11). Using $v_{t} = g (1 - β^{t})$ , find the exact number of steps $t$ needed for the velocity to reach 90% of its steady value.

Recall Solution 3.2

We need $1 - β^{t} \geq 0.9$ , i.e. $β^{t} \leq 0.1$ : $0. 9^{t} \leq 0.1 \Rightarrow t \geq \frac{l n 0.1}{l n 0.9} = \frac{- 2.3026}{- 0.1054} = 21.85.$ So $t = 22$ steps to reach 90% of cruising velocity. The " $1/ (1 - β) = 10$ " is the characteristic time constant (reaching $\approx 63%$ ), not the 90% mark — a useful nuance: the rule of thumb underestimates full recovery. Either way, resetting momentum costs you tens of steps of wasted crawling (see the shaded gap in the figure).

Level 4 — Synthesis

Exercise 4.1 (L4)

Design a checkpoint policy for a 7-day training run of a 2 GB model, given a disk budget of 30 GB. You want: (a) crash recovery, (b) the single best-validation model kept forever, (c) a few recent fallbacks in case validation loss spikes. Write the policy and verify it fits the budget.

Recall Solution 4.1

A clean three-slot policy:

latest.pt — overwritten every epoch (crash recovery). Cost: $2$ GB.
best.pt — overwritten only when validation improves (Strategy 1, save-best). Cost: $2$ GB.
Rolling window of last $K = 3$ periodic checkpoints (Strategy 2). Cost: $3 \times 2 = 6$ GB.

Total: $2 + 2 + 6 = 10 GB \leq 30 GB$ . ✅ Fits with 20 GB to spare (headroom for optimizer state, which for Adam roughly doubles per-checkpoint size — worth noting).

Why this covers all three needs: latest.pt = requirement (a); best.pt = (b), never lost even if later epochs overfit; the rolling window = (c), lets you reload from just before a spike. Connects to Learning Rate Schedules because a spike at a schedule boundary is exactly when you'd reach for a fallback.

Exercise 4.2 (L4)

An EarlyStopping(monitor='val_loss', patience=10, restore_best_weights=True) callback is running. Validation loss over epochs is:

epoch	30	31	32	33	...	42
val_loss	0.50	0.48	0.49	0.50	(no new best)	0.51

Epoch 31 is the best. After that, no epoch beats 0.48. At which epoch does training stop, and which weights are restored?

Recall Solution 4.2

Patience counts epochs since the last improvement. Best was epoch 31. Patience of 10 means training stops after 10 consecutive non-improving epochs: $stop epoch = 31 + 10 = 41.$ Training halts at the end of epoch 41 (epochs 32–41 = 10 non-improving epochs). Because restore_best_weights=True, the model's weights are rolled back to the epoch-31 checkpoint (val_loss 0.48) — you don't keep the worse epoch-41 weights.

Level 5 — Mastery

Exercise 5.1 (L5)

You saved a checkpoint from a model whose final layer was nn.Linear(256, 10) (10 classes). Now you want to reuse it on a new task with 37 classes, so your new model ends in nn.Linear(256, 37). What exactly happens if you call model.load_state_dict(checkpoint['model_state_dict']), and what's the correct load procedure? Name the technique.

Recall Solution 5.1

What happens: load_state_dict (strict by default) validates tensor shapes. The final layer weight in the checkpoint has shape $10 \times 256$ ; your new layer expects $37 \times 256$ . Mismatch → it raises a size-mismatch error for that layer. This is a feature: it catches architecture mismatches instead of silently loading garbage.

Correct procedure — load everything except the final layer:

Load with strict=False, or delete the mismatched keys from the state dict before loading, so the shared feature layers get their pretrained weights.
Leave the new nn.Linear(256, 37) randomly initialised (fresh head).
Fine-tune on the new task.

This is exactly Transfer Learning: keep the learned feature extractor, replace and retrain only the task-specific head.

Exercise 5.2 (L5)

For inference-only export, the parent note calls model.eval() before predicting and wraps the forward pass in torch.no_grad(). Explain, at the level of what computation changes, why each is needed — and give one concrete wrong prediction that skipping model.eval() could cause.

Recall Solution 5.2

model.eval() switches layers that behave differently in train vs. inference:
- Dropout stops randomly zeroing activations (in eval it passes everything through). Skipping this means predictions become random — the same input can give different outputs each call.
- BatchNorm stops using the current batch's statistics and instead uses the running averages accumulated during training. Skipping this makes a single-sample prediction depend on whatever else is in the batch — and for batch size 1 the batch variance is ill-defined, giving wildly wrong outputs.
- Concrete wrong prediction: classifying one image with dropout still active could label a clear "cat" as "dog" simply because half its features were randomly zeroed on that call.
torch.no_grad() tells PyTorch not to build the autograd computation graph. It changes no numbers — predictions are identical — but it saves memory and speeds up inference by skipping gradient bookkeeping you'll never use at deployment.

Rule: eval() changes outputs (correctness); no_grad() changes efficiency only.