Saving and loading models (checkpoints): Worked examples

This page is the "run every case" companion to the parent checkpoint note (its Hinglish twin). The parent told you what a checkpoint is and why the optimizer state matters. Here we do the opposite of theory: we hammer through every awkward situation a checkpoint can throw at you — resuming, exporting, momentum resets, disk-budget math, mismatched architectures — and we compute the actual numbers so you never get surprised at 3 a.m. when training crashes.

Before we start, one word we will use constantly, defined from zero:

The scenario matrix

Every checkpoint problem falls into one of these cells. The worked examples below are each tagged with the cell they cover, so by the end no cell is left dark.

Cell	Situation	What makes it tricky
A	Resume mid-training (crash)	Off-by-one on the epoch counter
B	Momentum reset (save weights only)	Velocity $v \to 0$ , slow rebuild
C	Momentum preserved (full state)	Training continues smoothly — the contrast case
D	Export for inference (no optimizer)	Smaller file, `eval()` mode, no state
E	Degenerate: zero checkpoints on disk	Fresh start, `start_epoch = 0`
F	Limiting case: $β \to 1$ or $β \to 0$	How long momentum takes to rebuild
G	Storage budget (word problem)	Multiply out GB, does it fit the disk?
H	Exam twist: architecture mismatch	Shapes disagree → load fails, why

Cell E — the degenerate "nothing on disk yet" case

Start here because it is the boundary: what happens when there is nothing to resume from?

Worked example E · First run, empty checkpoint folder

You launch training. The checkpoint directory is empty. What epoch does the loop start at, and what does the optimizer velocity equal?

Forecast: guess the two numbers before reading on.

Check the folder. sorted(glob('checkpoint_*.pt')) returns an empty list. Why this step? The resume logic only fires if files exist. No files → skip it.
Set the counter. start_epoch = 0. Why? We have completed zero epochs, so the first epoch we run is epoch $0$ .
Set the velocity. A fresh optimizer initialises $v_{0} = 0$ . Why? No gradient history exists yet — there is literally nothing to average.

Verify: start_epoch = 0 and $v_{0} = 0$ . Sanity: on the very first update the momentum formula $v_{1} = β v_{0} + (1 - β) \nabla$ collapses to $v_{1} = (1 - β) \nabla$ — pure gradient, exactly what vanilla SGD would do on step one. ✓

Cell A — resume after a crash (the off-by-one)

Worked example A · Crash at epoch 7, where do we restart?

Training is set for 20 epochs. It saves a checkpoint at the end of each epoch. The machine crashes during epoch 8, so the last successful save is checkpoint_0007.pt. At what epoch does the loop resume, and how many epochs still run?

Forecast: is it epoch 7 or 8? And is it 12 more epochs or 13?

Load the latest file. Sorting the zero-padded names puts 0007 last, so we load epoch-7 state. Why sort with zero-padding? "10" sorts before "7" alphabetically; "0010" vs "0007" fixes it.
Read the counter. The file stores 'epoch': 7 — the epoch we finished.
Add one. start_epoch = checkpoint['epoch'] + 1 = 8. Why +1? Epoch 7 is done. Repeating it would waste work and double-count that data. We resume at the next unfinished epoch.
Count remaining. The loop is range(start_epoch, 20) = range(8, 20). Why range? It stops before 20, giving epochs $8, 9, \dots, 19$ .

Verify: epochs run $= 20 - 8 = 12$ . Total completed = 8 (0–7) already + 12 remaining $= 20$ . ✓ No epoch is skipped or repeated.

Cell B vs Cell C — the momentum reset, computed

This is the heart of the parent note's warning. We turn it into real numbers.

Figure — Saving and loading models (checkpoints)

Recall the momentum recurrence (every symbol defined below the formula):

Worked example B · You saved weights only — how far below cruising speed is step one?

Assume the gradient is a constant $g = 1$ every step (a clean idealisation) and $β = 0.9$ . The model has been training long enough that velocity has reached its steady value. You resume, but you saved only weights, so $v$ resets to $0$ . How big is the very first velocity compared to the cruising velocity?

Forecast: guess what fraction of full speed the first step reaches.

Find cruising velocity $v_{\infty}$ . At steady state $v_{\infty} = β v_{\infty} + (1 - β) g$ . Why set $v_{t} = v_{t - 1}$ ? "Steady" means it stops changing. Solve: $v_{\infty} (1 - β) = (1 - β) g \Rightarrow v_{\infty} = g = 1$ .
First step after reset. $v_{1} = 0.9 \cdot 0 + 0.1 \cdot 1 = 0.1$ . Why? With $v_{0} = 0$ the recurrence keeps only the $(1 - β) g$ term.
Compare. $v_{1} / v_{\infty} = 0.1/1 = 0.1$ , i.e. 10% of cruising speed. Why does this matter? Your parameter update $θ_{t} = θ_{t - 1} - α v_{t}$ is ten times weaker on the first step — the model crawls while it rebuilds inertia.

Verify: step two gives $v_{2} = 0.9 (0.1) + 0.1 (1) = 0.19$ , still far from $1$ . The red curve in the figure climbs slowly — exactly the "lost inertia" the parent warned about. ✓

Cell F — the limiting cases of $β$

Worked example F · How many steps to rebuild momentum for extreme

$β$ ? The parent quoted a rebuild time of about $1/ (1 - β)$ steps. Compute it for three regimes and interpret the limits.

Forecast: which $β$ rebuilds instantly, which takes forever?

$β = 0.9$ (typical): $1/ (1 - 0.9) = 1/0.1 = 10$ steps. Why this formula? $(1 - β)$ is how much new gradient enters each step; its reciprocal is the characteristic timescale of the running average.
$β = 0.99$ (heavy momentum): $1/ (1 - 0.99) = 100$ steps. Why so long? Almost all weight goes to the old velocity, so a fresh $v = 0$ takes ages to fill up. Losing this state is expensive.
$β = 0$ (no momentum, plain SGD): $1/ (1 - 0) = 1$ step. Why? With no memory, there is nothing to rebuild — resetting costs nothing. This is the degenerate limit where saving optimizer state barely matters.

Verify: as $β \to 1$ , rebuild time $\to \infty$ (worst case to lose state); as $β \to 0$ , rebuild time $\to 1$ (nothing lost). The three values $10, 100, 1$ match. ✓

Cell G — the storage-budget word problem

Worked example G · Does 100 epochs fit on a 250 GB disk?

A model has 500 million parameters. Each parameter is a 32-bit float (4 bytes). You also save Adam optimizer state, which stores two extra numbers per parameter ( $m_{t}$ and $v_{t}$ ). You checkpoint every epoch for 100 epochs. Your disk has 250 GB free. Do you fit?

Forecast: guess the total GB before computing.

Bytes per parameter, all-in. weights + $m_{t}$ + $v_{t}$ = $3$ numbers $\times 4$ bytes $= 12$ bytes. Why $\times 3$ ? Adam keeps two moving averages on top of the weight itself.
One checkpoint size. $500, 000, 000 \times 12 = 6, 000, 000, 000$ bytes $= 6 GB$ (using $1 GB = 1 0^{9}$ bytes). Why $1 0^{9}$ ? Disk vendors count decimal GB; we match their convention.
Full run. $6 GB \times 100 epochs = 600 GB$ . Why multiply by frequency $\times$ duration? That is the parent's tradeoff equation: $Storage = size \times freq \times duration$ .
Compare to disk. $600 > 250$ . It does not fit.

Verify: $600 GB$ needed vs $250 GB$ available → overflow by $350 GB$ . ✓

The fix (Strategy 1 from the parent): save best-only, i.e. keep effectively $\sim 1$ checkpoint. $6 GB \leq 250 GB$ fits easily — a $100 \times$ reduction. The amber bar in the figure shows the collapse.

Cell D — export for inference (drop the optimizer)

Worked example D · How much smaller is an inference-only checkpoint?

Same 500 M-parameter model. For deployment you save weights only (no $m_{t}$ , no $v_{t}$ ) — see Model Deployment. By what factor does the file shrink versus the full training checkpoint?

Forecast: guess the shrink factor.

Inference file size. weights only $= 500, 000, 000 \times 4$ bytes $= 2 GB$ . Why drop the rest? At inference we never call .step(), so momentum buffers are dead weight.
Ratio to full checkpoint. $6 GB /2 GB = 3$ . Why exactly 3? Adam stored 3 numbers per parameter; keeping 1 gives a clean $3 \times$ cut.
Set eval mode. In code you call model.eval(). Why? It freezes dropout and switches BatchNorm to use stored running stats — otherwise predictions become random/unstable.

Verify: inference file $= 2 GB$ , exactly one-third of $6 GB$ . ✓ Smaller, faster to ship, and correct because gradients are never needed for a forward pass. (Reusing these weights to fine-tune later is the Transfer Learning scenario.)

Cell H — exam twist: architecture mismatch

Worked example H · Loading a 256-wide checkpoint into a 512-wide model

You saved a checkpoint where a Linear layer had shape $(784 \to 256)$ , so its weight tensor is $256 \times 784$ . Someone edited the model class to $(784 \to 512)$ , weight shape $512 \times 784$ . You call load_state_dict. What happens, and how many weights would have had to match?

Forecast: does it silently work, silently corrupt, or loudly fail?

Count checkpoint weights. $256 \times 784 = 200, 704$ .
Count model weights. $512 \times 784 = 401, 408$ . Why count? load_state_dict copies element-by-element; the tensors must be the same shape.
Compare. $200, 704 \neq = 401, 408$ , and $256 \neq = 512$ in the first dimension.
Result: load_state_dict raises a size-mismatch error for that layer key. Why loud, not silent? This is exactly the validation the parent praised: load_state_dict checks shapes and reports missing/unexpected/mismatched keys instead of corrupting memory.

Verify: ratio $401, 408/200, 704 = 2$ — the new layer wants twice as many weights, so no partial copy is possible; the load must fail. ✓ The fix: keep architecture and checkpoint in sync, or use strict=False only when you deliberately accept partial loading.

Recall

Recall Why resume at

epoch + 1, not epoch? Because the stored counter is the epoch we finished; the next unfinished one is one higher. Resuming at epoch + 1 avoids repeating work ::: it starts the loop at the first epoch never run.

Recall After saving weights only, what is the first velocity vs cruising velocity (β=0.9, constant g=1)?

First velocity ::: $0.1$ , which is $10%$ of the cruising velocity $1$ — momentum must rebuild over ~10 steps.

Recall 500M params, Adam, 4 bytes each, every epoch, 100 epochs — total storage?

$500 M \times 12 bytes \times 100 = 600 GB$ ::: (12 bytes = weight + $m_{t}$ + $v_{t}$ ).

Recall Why does inference-only export shrink an Adam checkpoint by 3×?

Adam stores 3 numbers per parameter (weight, $m_{t}$ , $v_{t}$ ) ::: inference keeps only the weight, so $6 GB \to 2 GB$ .

Recall Rebuild time for momentum when β = 0.99?

$1/ (1 - 0.99) = 100$ steps ::: heavier momentum is far more costly to lose.