Saving and Loading Models

Context

Saving a model writes its state dict to a file with torch.save. Loading builds the same network again in any later script, reads the file back with torch.load and copies the numbers in with load_state_dict. A checkpoint, Lesson 0.1.2's state saved to survive a crash, holds here what a run needs to go on, its parameters, its optimizer's state, the epoch it reached and the validation loss it scored there, and is the same call on a larger dictionary.

A model worth keeping is used by processes that did not train it: Section 4.5 saves its image network once and loads it in the next lesson instead of training it again, and Module 6's prediction service loads a model a training run logged. A run cut off at epoch 900 of 2,000 resumes at epoch 901 from a checkpoint instead of at epoch 1.

A network's parameters live in the memory of the process that fitted them, so Lesson 0.1.2's rule holds for them as it held for a counter: only what reached the disk before the process ended survives it. The quieter failure is the save itself. A loop that writes its file at every epoch without asking whether the model improved ends holding its last model, here a fifth worse on the validation rows than its best, and nothing raises. Lesson 4.4.2's early stopping kept the best epoch's parameters as a deep copy in memory, and this lesson puts that copy on disk.

Goal

• Read a network's state dict by name, shape and count, and name what it holds and what a loading script has to supply.
• Save the early-stopped network with torch.save, load it with torch.load(..., weights_only=True) and load_state_dict, and read the size mismatch and the missing and unexpected keys a load refuses, and the mismatch it lets through.
• Load the file in a second script, predict.py, reproduce the first script's weights and validation score, and score the test rows there once.
• Write a checkpoint of the epoch, the validation loss, the model and the optimizer, and resume a run from it to the validation losses an unbroken run prints.
• Show from a run that a file saved at every epoch keeps the last model and loses the best, and write the save that keeps the best.
• Explain why the default torch.load refuses a whole saved model, and what weights_only=False needs and risks.

Prerequisites

None of this is required, the lesson stands on its own. Both are here if you want to see a model saved and loaded before you do it yourself.

• YouTube - "PyTorch Tutorial 17 - Saving and Loading Models" by Patrick Loeber. Eighteen minutes on the three calls this lesson builds on, torch.save, torch.load and load_state_dict, the different ways of saving a model, and what changes when a GPU is involved. It was uploaded in April 2020, nearly five years before release 2.6 changed torch.load's default, so every load in it runs under the old one.
• Official docs - PyTorch Tutorials, "Saving and Loading Models", the section "What is a state_dict?". A state dict as a dictionary from each layer to its parameter tensors, with entries only for layers that hold learnable parameters and for registered buffers, batch normalization's running_mean among them, and an optimizer's own state dict of its state and its settings, both printed for a small image classifier.

Table of Contents

• state_dict Is the Model's Memory
• torch.save and torch.load
• Loading in a Second Process
• A Checkpoint With Its Epoch and Optimizer
• The Checkpoint That Overwrites the Best
• What Not to Save

Theory

Every block in this lesson runs in ~/dev/bootcamp on Lesson 4.3.1's table and the split Lesson 4.2.3 made, 1,725 fit rows, 575 validation rows and 575 test rows, and imports what earlier lessons saved: the frames and the six tensors of prep_net.py, loss_fn, network, validate, run and score from train_loop.py, early_stop from Lesson 4.4.2's stopping.py, and, in the second lab, the slice and fit from Lesson 4.4.1's slice_loop.py. The model is network(128), the 128-unit network of Lesson 4.3.1's width table, fitted on all 1,725 fit rows with optim.SGD at 0.1 and stopped early at a patience of 100. A Python block that opens with its imports is a script of its own, and a block that does not continues the script above it.

state_dict Is the Model's Memory

What. Lesson 4.2.2's state dict holds what a fit writes into a network. For network(128) it is four tensors under the names 0.weight, 0.bias, 2.weight and 2.bias, a module's position and a parameter's name as Lesson 4.3.1 read them, 2,561 numbers in all. The network itself is not in it: the classes, their order and the ReLU live in the code that builds the network. Why. A fitted model is the network its code builds plus the numbers the fit wrote into it, so a file that keeps the numbers keeps the model, provided some script can build the same network again. The file holds the numbers and the script holds the shape, and every refusal in the next subtitle is the two disagreeing. When. Whenever a model has to outlive the process that fitted it, which from this lesson on is every model worth keeping.

import torch
from torch import nn, optim
from prep_net import val_rows, X_fit, y_fit, X_val, y_val
from train_loop import network, validate, score
from stopping import early_stop

model = network(128)        # seeded 42 inside
state = model.state_dict()  # taken before a single step
for name, tensor in state.items():
    print(name, tuple(tensor.shape), tensor.dtype)
print(sum(t.numel() for t in state.values()))
print(f"{validate(model, X_val, y_val):.4f}")

0.weight (128, 18) torch.float32
0.bias (128,) torch.float32
2.weight (1, 128) torch.float32
2.bias (1,) torch.float32
2561
15.8030

The shapes say how wide each layer is, nothing in the dictionary says that the three modules are two nn.Linear layers with a ReLU between them, and 15.8030 is the seeded network on the validation rows before a single step. One kind of layer keeps more than parameters, the batch normalization of Lesson 4.4.2's Assignment:

norm = nn.BatchNorm1d(18)  # batch normalization over 18 columns
print([name for name, p in norm.named_parameters()])
print([name for name, b in norm.named_buffers()])
print(list(norm.state_dict()))

['weight', 'bias']
['running_mean', 'running_var', 'num_batches_tracked']
['weight', 'bias', 'running_mean', 'running_var', 'num_batches_tracked']

Two parameters and three buffers. A state dict holds every parameter and every buffer, which on the networks this lesson saves, with no buffer in them, is every parameter. Module 5's convolutional networks carry batch normalization, Lesson 5.1.2's subject, whose version for images keeps the same three buffers. The fit is Lesson 4.4.2's early_stop, which keeps a deep copy of the state dict at every new low and restores the best when it stops:

optimizer = optim.SGD(model.parameters(), lr=0.1)
# patience 100, at most 2,000 epochs, the best restored
stopped, best_epoch, best_loss = early_stop(model, optimizer, X_fit, y_fit, 100, 2000)
print(stopped, best_epoch, f"{best_loss:.4f}")
print(f"{validate(model, X_val, y_val):.4f}")
print(state["0.weight"][0, :3])
mae, r2 = score(model, X_val, val_rows)
print(f"{mae:,.0f}", round(r2, 3))

348 248 0.4412
0.4412
tensor([ 0.2293,  0.0613, -0.0976])
44,825 0.859

The loop stopped at epoch 348, a hundred epochs after its lowest validation loss, 0.4412 at epoch 248, the epoch and the loss Lesson 4.3.1's width table printed for this width, and restored epoch 248's numbers, which miss a validation listing by 44,825 on average. state, taken before the first step, prints epoch 248's weights: its tensors are the parameters' own memory, Lesson 4.2.3's live, and load_state_dict copies values into that memory rather than swapping in new tensors.

DON'T:

• DON'T expect a saved file to rebuild the network. It holds four names, four shapes and 2,561 numbers, and the script that loads it has to build network(128) first, with the code that built this one.
• DON'T call a state dict the model's parameters. nn.BatchNorm1d(18) lists two parameters and five state dict entries, so a copy made from parameters() carries two of its five tensors.

torch.save and torch.load

What. torch.save(model.state_dict(), "price_mlp.pt") is Lesson 4.1.1's torch.save handed the state dict, and torch.load("price_mlp.pt", weights_only=True) reads it back as a new dictionary of new tensors. weights_only=True, written out although it is the default, is the restriction behind the refusal of a NumPy array that Lesson 4.1.1 met, and the fourth subtitle reads what it lets through. fresh.load_state_dict(loaded) then copies the values into a network the loading script built. Why. load_state_dict matches the file to the network by name and by shape and refuses a mismatch with a message that names every key it could not place. It cannot check what the file does not hold, and the function between the layers is one of those things. When. Save once the epoch is chosen, here after early_stop restored epoch 248. Load at the top of every script that uses the model, into the network the fitting script's code builds, and call eval() before the first prediction.

torch.save(model.state_dict(), "price_mlp.pt")
loaded = torch.load("price_mlp.pt", weights_only=True)
print(type(loaded).__name__, list(loaded))
print(all(torch.equal(loaded[k], state[k]) for k in state))
fresh = network(128)  # the seeded start again
print(fresh.load_state_dict(loaded))
print(f"{validate(fresh, X_val, y_val):.4f}")

OrderedDict ['0.weight', '0.bias', '2.weight', '2.bias']
True
<All keys matched successfully>
0.4412

Every tensor comes back equal to the one saved, and the seeded start that scored 15.8030 scores 0.4412 once the numbers are in, with Lesson 4.2.3's report that every key matched. A network of another width has other shapes:

try:
    network(256).load_state_dict(loaded)
except RuntimeError as err:
    for line in str(err).splitlines():
        print(line.strip())

Error(s) in loading state_dict for Sequential:
size mismatch for 0.weight: copying a param with shape torch.Size([128, 18]) from checkpoint, the shape in current model is torch.Size([256, 18]).
size mismatch for 0.bias: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([256]).
size mismatch for 2.weight: copying a param with shape torch.Size([1, 128]) from checkpoint, the shape in current model is torch.Size([1, 256]).

One line per tensor whose shape differs, the file's shape first and the network's second, and 2.bias is not among them, because the output layer's bias is (1,) at every width. Each line of the message after the first opens with a tab, which strip() takes off. A layer added between the two changes the names instead:

torch.manual_seed(42)
dropped = nn.Sequential(nn.Linear(18, 128), nn.ReLU(), nn.Dropout(0.5), nn.Linear(128, 1))
print(list(dropped.state_dict()))
try:
    dropped.load_state_dict(loaded)
except RuntimeError as err:
    for line in str(err).splitlines():
        print(line.strip())
print(dropped.load_state_dict(loaded, strict=False))
print(f"{validate(dropped, X_val, y_val):.4f}")

['0.weight', '0.bias', '3.weight', '3.bias']
Error(s) in loading state_dict for Sequential:
Missing key(s) in state_dict: "3.weight", "3.bias".
Unexpected key(s) in state_dict: "2.weight", "2.bias".
_IncompatibleKeys(missing_keys=['3.weight', '3.bias'], unexpected_keys=['2.weight', '2.bias'])
15.3290

nn.Dropout(0.5) at position 2 moves the output layer to 3., as Lesson 4.4.2 found, so the network's 3. keys are missing from the file, and the file's 2. keys are unexpected in the network. strict=False, an argument that reports missing and unexpected keys instead of raising on them, loaded the first layer, left the output layer at its seeded draws and validated at 15.3290, next to the untrained 15.8030, with no error. The third case prints no message at all:

torch.manual_seed(42)
bent = nn.Sequential(nn.Linear(18, 128), nn.Tanh(), nn.Linear(128, 1))
print(bent.load_state_dict(loaded))
print(f"{validate(bent, X_val, y_val):.4f}")

<All keys matched successfully>
2.0960

nn.Tanh, the hyperbolic tangent of Lesson 4.3.1's Assignment, holds no parameter either, so this network has the file's four names and four shapes, and the report says every key matched. It computes with numbers fitted behind a ReLU and validates at 2.0960, nearly five times the 0.4412 the same numbers score behind the ReLU. The report checks names and shapes, and only the code that builds the network can check the rest.

The last call after a load is eval(), which Lesson 4.2.3 showed changes nothing on a network without dropout or batch normalization and which decides, on one with dropout, whether a prediction is a price or a draw. The lesson's network with nn.Dropout(0.5) after its ReLU, early-stopped and saved, is loaded into a second copy, seeded first so its dropout draws repeat on a rerun, and asked for the first validation listing's price, in thousands, three times in each mode:

torch.manual_seed(42)
trained = nn.Sequential(nn.Linear(18, 128), nn.ReLU(), nn.Dropout(0.5), nn.Linear(128, 1))
optimizer = optim.SGD(trained.parameters(), lr=0.1)
stopped, best_epoch, best_loss = early_stop(trained, optimizer, X_fit, y_fit, 100, 2000)
print(stopped, best_epoch, f"{best_loss:.4f}")
torch.save(trained.state_dict(), "dropout_mlp.pt")

torch.manual_seed(42)
served = nn.Sequential(nn.Linear(18, 128), nn.ReLU(), nn.Dropout(0.5), nn.Linear(128, 1))
print(served.load_state_dict(torch.load("dropout_mlp.pt", weights_only=True)))
print(served.training)
with torch.no_grad():
    print([round(served(X_val[:1]).item() * 100) for _ in range(3)])
    served.eval()  # evaluation mode, every unit kept
    print([round(served(X_val[:1]).item() * 100) for _ in range(3)])

515 415 0.4425
<All keys matched successfully>
True
[437, 403, 438]
[425, 425, 425]

The dropout network stopped at epoch 515 with its best, 0.4425, at epoch 415. The loaded copy reports training as True, because a new module starts in training mode and load_state_dict sets no mode, so its dropout layer switches off a fresh random share of the 128 hidden values at every call: one listing comes out at 437, 403 and 438 thousand, and after eval() at 425 thousand every time.

DON'T:

• DON'T trust <All keys matched successfully> as proof that the network is the one that was fitted. The Tanh network matched all four keys and validates at 2.0960, where the same numbers behind a ReLU validate at 0.4412.
• DON'T pass strict=False to quiet a refusal. It loaded the first layer, left the output layer at its seeded draws, validated at 15.3290, and raised nothing.
• DON'T predict with a loaded network before eval(). With dropout in it, one listing came out at three prices from 403 to 438 thousand.

Loading in a Second Process

What. predict.py is a script that never fits. It builds network(128) with train_loop's code, loads price_mlp.pt, switches to evaluation mode and scores. Its import prep_net rebuilds the split, the four means and standard deviations and the 14 categories from listings.csv in the new process, so the network reads the same 18 columns it was fitted on, and none of that comes from the file. Why. A second process that prints the first one's weights and scores proves that the file carried the model, since the fitting script exited before this one started. The test rows have waited for a model chosen on the validation rows, and this is that model. When. Whenever the model is used, from a fresh process, and the test rows are scored here, once.

Save it as predict.py and run it once the fitting script has finished:

import torch
from prep_net import val_rows, test_rows, X_val, X_test
from train_loop import network, score

model = network(128)  # the shape, from code
print(model.load_state_dict(torch.load("price_mlp.pt", weights_only=True)))
model.eval()          # before any prediction
print(model.state_dict()["0.weight"][0, :3])
mae, r2 = score(model, X_val, val_rows)
print(f"{mae:,.0f}", round(r2, 3))
mae, r2 = score(model, X_test, test_rows)
print(f"{mae:,.0f}", round(r2, 3))

(bootcamp) goku@machine:~/dev/bootcamp$ python3 predict.py
<All keys matched successfully>
tensor([ 0.2293,  0.0613, -0.0976])
44,825 0.859
45,862 0.843

The first unit's first three weights are the first script's, and the validation rows score 44,825 and 0.859 as they did there, so the 2,561 numbers crossed from one process to the other through the file. score sets the modes itself, so model.eval() changes nothing in these lines, and it stays as the habit the dropout network asked for, the line before any direct call of model(X). The test rows, scored once, put the network's mean miss at 45,862 and its R² at 0.843, the one score in this lesson that stands for listings no choice was made on.

DON'T:

• DON'T score the test rows in the fitting script. They waited for the model the validation rows chose, and they are scored in predict.py, once.
• DON'T assume the file carries the table's statistics. Its four keys are the network's, and mu, sd and enc came from import prep_net running the split again on listings.csv.

A Checkpoint With Its Epoch and Optimizer

What. A checkpoint meant for resuming is one dictionary handed to torch.save, with the keys epoch, val_loss, model for the network's state dict and optimizer for optimizer.state_dict(). The optimizer's dictionary has two keys of its own: state, what Lesson 4.3.4 found each optimizer keeping between steps, held under each parameter's position from 0 to 3, and param_groups, the settings, with the same positions listed under params. Why. The model's state dict is enough to predict and not enough to resume: a run restarted from it alone starts the optimizer's state from nothing and counts its epochs from 1, where a resumed run is worth having only when it takes exactly the steps the interrupted run would have taken. The epoch and the validation loss also say which model the file holds. When. Every run long enough to be interrupted, every few epochs or at every new low.

The lesson's network, run for 248 epochs as if the run were cut there:

import pickle
import numpy as np
import torch
from torch import optim
from train_loop import network, run

model = network(128)
optimizer = optim.SGD(model.parameters(), lr=0.1)
train_losses, val_losses = run(model, optimizer, 248)
saved = optimizer.state_dict()
print(list(saved), saved["state"])
print(saved["param_groups"][0]["lr"], saved["param_groups"][0]["params"])
checkpoint = {
    "epoch": 248,
    "val_loss": val_losses[-1],           # a Python float
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),  # its state and its settings
}
torch.save(checkpoint, "checkpoint.pt")

['state', 'param_groups'] {}
0.1 [0, 1, 2, 3]

Plain SGD keeps nothing between steps, Lesson 4.3.4's finding, so its state is empty, and param_groups holds the rate and the positions of the four tensors it updates, in the order model.parameters() gives them. The file names the parameters by position alone, so an optimizer is built over the new network's parameters before its state is loaded. Loaded and resumed for 100 epochs, beside a run of 348 that was never cut:

back = torch.load("checkpoint.pt", weights_only=True)
print(list(back), back["epoch"], f"{back['val_loss']:.4f}")
resumed = network(128)
resumed.load_state_dict(back["model"])
# built over the new network first, then loaded
resumed_optimizer = optim.SGD(resumed.parameters(), lr=0.1)
resumed_optimizer.load_state_dict(back["optimizer"])
more_train, more_val = run(resumed, resumed_optimizer, 100)

unbroken = network(128)
whole_train, whole_val = run(unbroken, optim.SGD(unbroken.parameters(), lr=0.1), 348)
print(more_val == whole_val[248:], f"{more_val[-1]:.4f}")
print(all(torch.equal(a, b) for a, b in zip(resumed.parameters(), unbroken.parameters())))

['epoch', 'val_loss', 'model', 'optimizer'] 248 0.4412
True 0.4444
True

run counts its epochs from 1, so the resumed run's 100 are the unbroken run's epochs 249 to 348, and every validation loss equals the unbroken run's, down to the parameters at the end. Its last, 0.4444, is epoch 348's, the epoch at which early_stop stopped. SGD at 0.1 resumes just as exactly with its empty state never loaded, because there is nothing in it to lose. Momentum keeps a buffer:

# a run cut at epoch 100, its buffer saved with it
model = network(128)
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
run(model, optimizer, 100)
print(list(optimizer.state_dict()["state"][0]))
torch.save({"epoch": 100, "model": model.state_dict(), "optimizer": optimizer.state_dict()}, "momentum.pt")
halfway = torch.load("momentum.pt", weights_only=True)

curves = []
# True loads the saved buffer, False starts without it
for keep in (True, False):
    model = network(128)
    model.load_state_dict(halfway["model"])
    optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
    if keep:
        optimizer.load_state_dict(halfway["optimizer"])
    curves.append(run(model, optimizer, 100)[1])
unbroken = network(128)
whole_val = run(unbroken, optim.SGD(unbroken.parameters(), lr=0.01, momentum=0.9), 200)[1]
print([curve == whole_val[100:] for curve in curves])
print([f"{curve[0]:.4f}" for curve in curves], f"{whole_val[100]:.4f}")
print([f"{curve[-1]:.4f}" for curve in curves], f"{whole_val[-1]:.4f}")

['momentum_buffer']
[True, False]
['0.4669', '0.4677'] 0.4669
['0.4566', '0.4571'] 0.4566

Momentum at 0.9 keeps one momentum_buffer per parameter, Lesson 4.3.4's, and the buffer crossed the file with the model: resumed with it, the run is the unbroken run. Resumed without it, a fresh optimizer starts its buffer as the first gradient it sees, Lesson 4.3.4's rule, so the first resumed epoch lands at 0.4677 where the run it continues lands at 0.4669, and the run ends at 0.4571 instead of 0.4566. The gap is small on this network, and from the first resumed step it is another run.

The two checkpoints passed weights_only=True holding an int, a Python float, dictionaries of tensors and the optimizer's settings, numbers, lists, False and None among them. What it refuses, with pickle.UnpicklingError (an error of pickle, Python's module that turns objects into bytes and back), is a value whose class is not on its list, and a NumPy number is the one a training script makes without noticing:

lowest = np.min(whole_val)
print(type(lowest).__name__, type(min(whole_val)).__name__)
torch.save({"epoch": 200, "val_loss": lowest}, "numpy_loss.pt")
try:
    torch.load("numpy_loss.pt", weights_only=True)
except pickle.UnpicklingError as err:
    print(type(err).__name__)
    for line in str(err).splitlines():
        if "GLOBAL" in line:
            print(line.strip())
torch.save({"epoch": 200, "val_loss": float(lowest)}, "numpy_loss.pt")
print(torch.load("numpy_loss.pt", weights_only=True)["epoch"])

float64 float
UnpicklingError
WeightsUnpickler error: Unsupported global: GLOBAL numpy._core.multiarray.scalar was not an allowed global by default. Please use `torch.serialization.add_safe_globals([numpy._core.multiarray.scalar])` or the `torch.serialization.safe_globals([numpy._core.multiarray.scalar])` context manager to allowlist this global if you trust this class/function.
200

np.min returns a NumPy float64 where Python's own min on the same list returns a float, and the whole file is refused over that one value. The message names a global (a class or function a file names by its module path, for the loader to import and call), numpy._core.multiarray.scalar, the function NumPy rebuilds such a number with, which the loader will not call by default. float(...) around the value makes it Python's, and the file loads.

DON'T:

• DON'T leave the optimizer out of a checkpoint meant for resuming. Plain SGD lost nothing because it keeps nothing, and momentum resumed without its buffer landed at 0.4677 where the run it continued landed at 0.4669.
• DON'T store a NumPy number in a checkpoint. np.min returned a float64, and weights_only=True refused the whole file over it.

The Checkpoint That Overwrites the Best

What. Two saves in one loop: latest.pt, written at every epoch with no question asked, and best.pt, written only when the validation loss sets a new low. The first is the save of Lesson 0.1.2's crash_safe.py, which overwrote checkpoint.txt at every epoch under the comment "overwrite the checkpoint every epoch (latest wins)". Why. A save replaces the file it writes. Past the epoch of the lowest validation loss, every unconditional save writes a model the validation rows score worse than the best, and nothing raises: the file is valid, loads, matches every key and scores worse. In Lesson 0.1.2 the latest value of a counter was the one worth keeping, and in a fit the latest model stops being the best one the moment the validation loss turns. When. Every loop that saves inside its epochs. The file you will load for predictions is written under a condition, and an unconditional file, if you keep one, exists for resuming and is named for it.

import torch
from torch import optim
from prep_net import val_rows, X_fit, y_fit, X_val, y_val
from train_loop import loss_fn, network, validate, score

model = network(128)
optimizer = optim.SGD(model.parameters(), lr=0.1)
best_loss, best_epoch, writes = float("inf"), 0, 0
for epoch in range(1, 2001):
    loss = loss_fn(model(X_fit), y_fit)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    val_loss = validate(model, X_val, y_val)
    # every epoch, no question asked
    torch.save(model.state_dict(), "latest.pt")
    if val_loss < best_loss:
        best_loss, best_epoch, writes = val_loss, epoch, writes + 1
        # only on a new low
        torch.save(model.state_dict(), "best.pt")
print(epoch, f"{val_loss:.4f}", best_epoch, f"{best_loss:.4f}", writes)

2000 0.5279 248 0.4412 114

scores = {}
for name in ("latest.pt", "best.pt"):
    kept = network(128)
    print(kept.load_state_dict(torch.load(name, weights_only=True)))
    scores[name] = validate(kept, X_val, y_val)
    mae, r2 = score(kept, X_val, val_rows)
    print(name, f"{scores[name]:.4f}", f"{mae:,.0f}", round(r2, 3))
print(round(scores["latest.pt"] / scores["best.pt"], 1))

<All keys matched successfully>
latest.pt 0.5279 49,659 0.832
<All keys matched successfully>
best.pt 0.4412 44,825 0.859
1.2

best.pt was written 114 times, the last at epoch 248, and latest.pt 2,000 times, the 1,752 from epoch 249 on each holding a model the validation rows score worse than the one in best.pt. Both files load with every key matched, and the latest misses a validation listing by 49,659 on average where the best misses by 44,825, its validation loss 1.2 times the best one's. On one or two threads the same script prints the best file's numbers to the digit and 0.5261 and 49,573 for the latest: the end of this run moves with the thread count and its lowest point does not, as Lesson 4.4.1 measured on the slice.

DON'T:

• DON'T load a file saved at every epoch for predictions. latest.pt validates at 0.5279 against best.pt's 0.4412, and it loaded with every key matched.
• DON'T judge a save by whether its file loads. Both files loaded without a word, and only a validation pass told them apart.

What Not to Save

What. torch.save(model, "whole.pt") writes the module object itself, as joblib.dump wrote the fitted pipeline of Lesson 3.2.3's Assignment: its class, named by module and name, and everything the object holds. The default torch.load refuses such a file, and weights_only=False rebuilds the object by importing its class, which works only where that class imports under the same name. Why. A state dict is names and tensors, which the restricted load reads. A whole object is instructions to import classes and call them, and a file can carry instructions to call anything, the risk Lesson 3.2.3 named for joblib.load, so PyTorch's default refuses it. A model of your own class also ties the file to the script that defined the class. When. Never for a model another script will load. The state dict is what travels, and the class lives in code.

import pickle
import torch
from prep_net import X_val, y_val
from train_loop import network, validate

model = network(128)
model.load_state_dict(torch.load("price_mlp.pt", weights_only=True))
torch.save(model, "whole.pt")  # the object, not its state dict
try:
    torch.load("whole.pt")     # the default, weights_only=True
except pickle.UnpicklingError as err:
    print(type(err).__name__)
    for line in str(err).splitlines():
        if "GLOBAL" in line or "arbitrary" in line:
            print(line.strip())
whole = torch.load("whole.pt", weights_only=False)
print(type(whole).__name__, f"{validate(whole, X_val, y_val):.4f}")

UnpicklingError
(1) In PyTorch 2.6, we changed the default value of the `weights_only` argument in `torch.load` from `False` to `True`. Re-running `torch.load` with `weights_only` set to `False` will likely succeed, but it can result in arbitrary code execution. Do it only if you got the file from a trusted source.
WeightsUnpickler error: Unsupported global: GLOBAL torch.nn.modules.container.Sequential was not an allowed global by default. Please use `torch.serialization.add_safe_globals([torch.nn.modules.container.Sequential])` or the `torch.serialization.safe_globals([torch.nn.modules.container.Sequential])` context manager to allowlist this global if you trust this class/function.
Sequential 0.4412

The two lines kept from the message say what weights_only=False risks and name the global the file asks for, torch.nn.modules.container.Sequential. With weights_only=False the file loads and validates at 0.4412, because torch.nn imports wherever PyTorch is installed. A class of your own is another matter. Save this as save_whole.py:

import torch
from torch import nn
from train_loop import network

class PriceNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = network(128)

    def forward(self, X):
        return self.layers(X)

net = PriceNet()
net.layers.load_state_dict(torch.load("price_mlp.pt", weights_only=True))
torch.save(net, "price_net.pt")
print(list(net.state_dict()))

['layers.0.weight', 'layers.0.bias', 'layers.2.weight', 'layers.2.bias']

PriceNet is Lesson 4.2.2's kind of nn.Module subclass, holding network(128) as self.layers, and its state dict's names follow the attributes, layers. in front of the Sequential's own, which is why the price file is loaded into net.layers. Then load the whole file from a second script, load_whole.py:

import torch

try:
    torch.load("price_net.pt", weights_only=False)
except AttributeError as err:
    print(type(err).__name__, str(err).split(" on ")[0])

AttributeError Can't get attribute 'PriceNet'

The message goes on to name the module it searched, the main script of the run. PriceNet was defined in save_whole.py, which ran as the main script, so the file asks the loader for the main script's PriceNet, and load_whole.py, the main script when it runs, defines none. The state dict needed nothing but network(128), code train_loop.py already holds.

DON'T:

• DON'T torch.save(model) a model another script will load. The default torch.load refused it, and weights_only=False could not find PriceNet in a script that does not define it.
• DON'T pass weights_only=False to a file you did not write. The refusal itself says why: the load can result in arbitrary code execution.

Lab Exercises

Lab Exercise 1: The 32-Unit Network in a File

In save32.py, build Lesson 4.3.1's 32-unit network with network(), take its state dict, and print each name with its shape, then the total count. Fit it with early_stop at a patience of 100 within 2,000 epochs at optim.SGD 0.1, print the stopping epoch, the best epoch and the best validation loss, then the first three weights of the first unit from the dictionary taken before the fit, and save the state dict as price_mlp32.pt.

Load the file into a fresh network() and print the report and the validation loss. Try loading the Theory's price_mlp.pt into the same network and print the message line by line. Then write Wrapped, an nn.Module whose __init__ keeps network() as self.layers and whose forward returns self.layers(X), print its state dict's names, try loading price_mlp32.pt into it and print the message, and load the file into wrapped.layers instead, printing the report and the validation loss.

Finally write predict32.py, which loads price_mlp32.pt into a fresh network(), switches it to evaluation mode, and prints the report, the first three weights of the first unit, and the validation MAE and R², without touching the test rows.

Lab Exercise 2: Checkpoints on the Memorized Slice

In slice_checkpoints.py, fit network(512) at optim.SGD 0.1 for 3,000 epochs on X_slice and y_slice, the 200 fit rows slice_loop.py exports, keeping every epoch's validation loss in a list. At every epoch save a checkpoint of the epoch, the validation loss, the model and the optimizer as slice_latest.pt, and the same checkpoint as slice_best.pt whenever the validation loss sets a new low.

After the loop, for each file print its name, its stored epoch and validation loss, the validation loss of a fresh network(512) loaded from it and that network's validation MAE in whole units, then the ratio of the two recomputed losses. Resume a fresh network and optimizer from slice_best.pt for 100 epochs with slice_loop's fit, and print whether their validation losses equal the original run's from the epoch after the best, and the last of them. Finish by saving the resumed network whole, loading the file with the default torch.load, and printing the exception's class name and the line of its message that names the class it refused.

import torch
from torch import nn, optim
from prep_net import X_fit, y_fit, X_val, y_val
from train_loop import network, validate
from stopping import early_stop

model = network()           # width 32, the default
state = model.state_dict()  # before the fit
for name, tensor in state.items():
    print(name, tuple(tensor.shape))
print(sum(t.numel() for t in state.values()))
optimizer = optim.SGD(model.parameters(), lr=0.1)
stopped, best_epoch, best_loss = early_stop(model, optimizer, X_fit, y_fit, 100, 2000)
print(stopped, best_epoch, f"{best_loss:.4f}")
print(state["0.weight"][0, :3])
torch.save(model.state_dict(), "price_mlp32.pt")

0.weight (32, 18)
0.bias (32,)
2.weight (1, 32)
2.bias (1,)
641
466 366 0.4424
tensor([ 0.7399, -0.0632, -0.1077])

fresh = network()
print(fresh.load_state_dict(torch.load("price_mlp32.pt", weights_only=True)))
print(f"{validate(fresh, X_val, y_val):.4f}")
try:
    fresh.load_state_dict(torch.load("price_mlp.pt", weights_only=True))
except RuntimeError as err:
    for line in str(err).splitlines():
        print(line.strip())

<All keys matched successfully>
0.4424
Error(s) in loading state_dict for Sequential:
size mismatch for 0.weight: copying a param with shape torch.Size([128, 18]) from checkpoint, the shape in current model is torch.Size([32, 18]).
size mismatch for 0.bias: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([32]).
size mismatch for 2.weight: copying a param with shape torch.Size([1, 128]) from checkpoint, the shape in current model is torch.Size([1, 32]).

class Wrapped(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = network()

    def forward(self, X):
        return self.layers(X)

wrapped = Wrapped()
print(list(wrapped.state_dict()))
try:
    wrapped.load_state_dict(torch.load("price_mlp32.pt", weights_only=True))
except RuntimeError as err:
    for line in str(err).splitlines():
        print(line.strip())
print(wrapped.layers.load_state_dict(torch.load("price_mlp32.pt", weights_only=True)))
print(f"{validate(wrapped, X_val, y_val):.4f}")

['layers.0.weight', 'layers.0.bias', 'layers.2.weight', 'layers.2.bias']
Error(s) in loading state_dict for Wrapped:
Missing key(s) in state_dict: "layers.0.weight", "layers.0.bias", "layers.2.weight", "layers.2.bias".
Unexpected key(s) in state_dict: "0.weight", "0.bias", "2.weight", "2.bias".
<All keys matched successfully>
0.4424

import torch
from prep_net import val_rows, X_val
from train_loop import network, score

model = network()
print(model.load_state_dict(torch.load("price_mlp32.pt", weights_only=True)))
model.eval()
print(model.state_dict()["0.weight"][0, :3])
mae, r2 = score(model, X_val, val_rows)
print(f"{mae:,.0f}", round(r2, 3))

(bootcamp) goku@machine:~/dev/bootcamp$ python3 predict32.py
<All keys matched successfully>
tensor([ 0.7399, -0.0632, -0.1077])
45,611 0.859

import pickle
import torch
from torch import optim
from prep_net import val_rows, X_val, y_val
from train_loop import loss_fn, network, validate, score
from slice_loop import X_slice, y_slice, fit

model = network(512)
optimizer = optim.SGD(model.parameters(), lr=0.1)
best_loss, val_losses = float("inf"), []
for epoch in range(1, 3001):
    loss = loss_fn(model(X_slice), y_slice)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    val_losses.append(validate(model, X_val, y_val))
    checkpoint = {
        "epoch": epoch,
        "val_loss": val_losses[-1],
        "model": model.state_dict(),
        "optimizer": optimizer.state_dict(),
    }
    torch.save(checkpoint, "slice_latest.pt")
    if val_losses[-1] < best_loss:
        best_loss = val_losses[-1]
        torch.save(checkpoint, "slice_best.pt")

scores = {}
for name in ("slice_latest.pt", "slice_best.pt"):
    saved = torch.load(name, weights_only=True)
    kept = network(512)
    kept.load_state_dict(saved["model"])
    scores[name] = validate(kept, X_val, y_val)
    mae, r2 = score(kept, X_val, val_rows)
    print(name, saved["epoch"], f"{saved['val_loss']:.4f}", f"{scores[name]:.4f}", f"{mae:,.0f}")
print(round(scores["slice_latest.pt"] / scores["slice_best.pt"], 1))

slice_latest.pt 3000 3.3360 3.3360 96,693
slice_best.pt 38 0.7764 0.7764 63,281
4.3

best = torch.load("slice_best.pt", weights_only=True)
resumed = network(512)
resumed.load_state_dict(best["model"])
resumed_optimizer = optim.SGD(resumed.parameters(), lr=0.1)
resumed_optimizer.load_state_dict(best["optimizer"])
again = fit(resumed, resumed_optimizer, X_slice, y_slice, 100)[1]
start = best["epoch"]
print(again == val_losses[start:start + 100], f"{again[-1]:.4f}")

torch.save(resumed, "slice_whole.pt")
try:
    torch.load("slice_whole.pt")
except pickle.UnpicklingError as err:
    print(type(err).__name__)
    for line in str(err).splitlines():
        if "GLOBAL" in line:
            print(line.strip())

True 1.0194
UnpicklingError
WeightsUnpickler error: Unsupported global: GLOBAL torch.nn.modules.container.Sequential was not an allowed global by default. Please use `torch.serialization.add_safe_globals([torch.nn.modules.container.Sequential])` or the `torch.serialization.safe_globals([torch.nn.modules.container.Sequential])` context manager to allowlist this global if you trust this class/function.

Theory Assessment

Q1: What the File Could Not Check

price_mlp.pt holds four keys. Say what each names and what its value is, then what each of the Theory's three loads into a network other than network(128) printed and why: network(256), the network with nn.Dropout(0.5) at position 2, and the network with nn.Tanh in place of the ReLU. Say which of the three no reader of the report could have caught, and why no check that reads the file could have caught it either.

Hint: List what a state dict holds beside the numbers, then list what differs between network(128) and each of the three networks. A check can compare only what both sides hold.

Q2: Two Processes and One File

predict.py printed tensor([ 0.2293, 0.0613, -0.0976]) and 44,825 0.859, the fitting script's numbers. Say what crossed from the first process to the second and how, what predict.py rebuilt for itself and from what, what model.eval() did on this network and on the dropout network, and why the test rows were scored in predict.py and nowhere else in the lesson.

Hint: The file holds four keys, so everything else the second process used came from its own code and its own data. For the dropout network, ask which mode a module is in when it is built, and whether load_state_dict changes it.

Q3: Resuming, and the Overwrite

The SGD run resumed from the epoch-248 checkpoint matched the unbroken run at every epoch, the momentum run resumed without its optimizer state did not, and after 2,000 epochs latest.pt validated at 0.5279 and best.pt at 0.4412. Say what a checkpoint holds beyond the model's state dict and why, why the SGD resume matches with or without the optimizer's state while the momentum resume needed it, and what each file held after 2,000 epochs and why both loaded with every key matched.

Hint: Picture optimizer.state_dict()["state"] after a step for each optimizer, and what a freshly built optimizer holds. For the two files, ask which epoch last wrote each, and what load_state_dict compares.

Q4: The Best Model, Saved Late

A teammate wants the first subtitle's early stopping without stopping.py, writes the loop out, and saves the best parameters once, after it:

import torch
from torch import optim
from prep_net import X_fit, y_fit, X_val, y_val
from train_loop import loss_fn, network, validate

model = network(128)
optimizer = optim.SGD(model.parameters(), lr=0.1)
best_loss, best_epoch = float("inf"), 0
for epoch in range(1, 2001):
    loss = loss_fn(model(X_fit), y_fit)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    val_loss = validate(model, X_val, y_val)
    if val_loss < best_loss:
        best_loss, best_epoch = val_loss, epoch
        best_state = model.state_dict()
    elif epoch - best_epoch >= 100:
        break
torch.save(best_state, "best_model.pt")
print(epoch, best_epoch, f"{best_loss:.4f}")

check = network(128)
check.load_state_dict(torch.load("best_model.pt", weights_only=True))
print(f"{validate(check, X_val, y_val):.4f}")

348 248 0.4412
0.4444

The teammate reads the two lines as proof that torch.save loses precision: the best epoch scored 0.4412, the file scores 0.4444, so float32 rounding on the way to disk cost 0.0032, and the fix is to save the parameters in float64. Say which model the file holds and why, what is wrong with the precision claim and which run of this lesson shows it, and what one change to the loop makes the file hold epoch 248.

Hint: Ask what model.state_dict() hands back at epoch 248 and what the object best_state names holds by epoch 348. Then ask what torch.equal said about a loaded tensor and the one saved, and where 0.4444 appeared before.

Assignment

Assignment: Resume a Mini-Batch Adam Run, Shuffle Included

Section 4.5 trains with Adam on mini-batches, and a run of that kind carries more state than this lesson's checkpoints held. Lesson 4.3.4's Adam keeps a step count and two running averages for every parameter, and Lesson 4.3.3's loader draws each epoch's order from its generator, which moves on with every epoch the loader runs, so a loader handed a freshly seeded generator replays the first epochs' orders. Two methods of torch.Generator that this lesson did not use carry the generator's position: get_state() returns it as a torch.uint8 tensor, which weights_only=True reads back like any other tensor, and set_state(state) puts a generator back at that position.

Write resume.py. Fit network(128) with optim.Adam at 0.001 on a DataLoader over TensorDataset(X_fit, y_fit) with batch_size=64, shuffle=True and a generator seeded 42, for 40 epochs without a break, keeping every epoch's validation loss and a copy of the final parameters, and print the validation loss at epochs 20 and 40. Then run the same fit from a fresh network, optimizer and generator for 20 epochs, save a checkpoint of the epoch, the validation loss, the model, the optimizer and the generator's state as resume.pt, load it, and print whether those 20 validation losses equal the unbroken run's first 20, the stored epoch and loss, the generator state's dtype and shape, and Adam's step count for the first parameter. Resume three times for the last 20 epochs, each from a fresh network(128) loaded from the file: with the saved optimizer and the saved generator, with a freshly seeded generator in place of the saved one, and with a fresh optimizer in place of the saved one. For each, print the two choices, whether its 20 validation losses equal the unbroken run's last 20, its final validation loss, and the largest absolute difference between its final parameters and the unbroken run's. End with a comment saying what each of the three states carries and which of them the fourth subtitle's full-batch SGD run needed.

Hints:

• Write the epoch loop once as train(model, optimizer, generator, epochs), which builds its loader from the generator it is handed, so the three resumes differ only in what they hand it. Adam's count sits at optimizer.state_dict()["state"][0]["step"], one step per batch.
• [p.detach().clone() for p in model.parameters()] keeps the unbroken run's final numbers apart from the network's memory, and (p - q).abs().max().item() is one pair of tensors' largest difference.

import torch
from torch import optim
from torch.utils.data import DataLoader, TensorDataset
from prep_net import X_fit, y_fit, X_val, y_val
from train_loop import loss_fn, network, validate

dataset = TensorDataset(X_fit, y_fit)

def train(model, optimizer, generator, epochs):
    loader = DataLoader(dataset, batch_size=64, shuffle=True, generator=generator)
    val_losses = []
    for epoch in range(epochs):
        for xb, yb in loader:
            loss = loss_fn(model(xb), yb)
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
        val_losses.append(validate(model, X_val, y_val))
    return val_losses

model = network(128)
optimizer = optim.Adam(model.parameters(), lr=1e-3)
unbroken = train(model, optimizer, torch.Generator().manual_seed(42), 40)
final = [p.detach().clone() for p in model.parameters()]
print(f"{unbroken[19]:.4f}", f"{unbroken[-1]:.4f}")

model = network(128)
optimizer = optim.Adam(model.parameters(), lr=1e-3)
generator = torch.Generator().manual_seed(42)
first = train(model, optimizer, generator, 20)
torch.save({
    "epoch": 20,
    "val_loss": first[-1],
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
    "generator": generator.get_state(),  # its place in the stream
}, "resume.pt")
saved = torch.load("resume.pt", weights_only=True)
print(first == unbroken[:20], saved["epoch"], f"{saved['val_loss']:.4f}")
print(saved["generator"].dtype, saved["generator"].shape, saved["optimizer"]["state"][0]["step"])

for keep_optimizer, keep_generator in ((True, True), (True, False), (False, True)):
    model = network(128)
    model.load_state_dict(saved["model"])
    optimizer = optim.Adam(model.parameters(), lr=1e-3)
    if keep_optimizer:
        optimizer.load_state_dict(saved["optimizer"])
    generator = torch.Generator().manual_seed(42)
    if keep_generator:
        generator.set_state(saved["generator"])
    rest = train(model, optimizer, generator, 20)
    gap = max((p - q).abs().max().item() for p, q in zip(model.parameters(), final))
    print(keep_optimizer, keep_generator, rest == unbroken[20:], f"{rest[-1]:.4f}", f"{gap:.4f}")
# The model carries the 2,561 parameters, the optimizer carries Adam's
# step count and its two averages for each of them, and the generator
# carries the place in the stream that orders every epoch's batches.
# The fourth subtitle's full-batch SGD at 0.1 needed the model alone:
# its optimizer kept nothing and its loop drew no order.

0.4612 0.4379
True 20 0.4612
torch.uint8 torch.Size([5056]) tensor(540.)
True True True 0.4379 0.0000
True False False 0.4364 0.2128
False True False 0.4340 0.1956

Additional Resources

Optional, and nothing later depends on either. Worth opening if you want to know what saving a slice of a tensor costs, or how a file written on a GPU loads on a processor.

• Official docs - PyTorch 2.14, "Serialization semantics", the section "Saving and loading tensors preserves views". Tensors that share storage still share it after a save and a load, a five-value slice of a 999-value tensor saves all 999 values with it, and .clone() before the save writes only the five, at the cost of the view.
• Official docs - PyTorch 2.14, torch.load. The map_location argument, which remaps where each saved tensor's storage lands, the page's note that a file holding GPU tensors loads them onto the GPU unless map_location='cpu' is passed before load_state_dict, and mmap, which maps the file instead of reading every tensor into memory.

Summary

A trained network turned out to be two things that travel separately. The numbers, 2,561 under four names, go into a file, and the network, its classes, their order and the function between its layers, stays in code that builds it the same way every time. Saving writes the first, loading runs the second and pours the first back in, and every refusal load_state_dict raised was the two disagreeing about a name or a shape. The one disagreement nothing caught was the function between the layers, which the file never held.

The second process made the split concrete. A script that never trained printed the first one's weights and validation score, rebuilt the table's statistics from the data rather than from the file, and scored the test rows once. A checkpoint carried the same idea from a model to a run, with the optimizer's state and the epoch beside the parameters, so that the steps after an interruption are the steps the run would have taken.

The failure worth carrying forward is the quiet one. A file saved at every epoch is always valid and always loads, and after the lowest validation loss it always holds a worse model than one the run already passed through. Lesson 0.1.2's counter wanted its latest state, and a fit wants its best, so the save has to ask. Section 4.5 takes the loop, the two switches and the file to images.

Cheat Sheet

Part A - Reference Table

Part B - Compact Code Block

state = model.state_dict()         # names to tensors, no class
torch.save(state, "price_mlp.pt")  # after the epoch is chosen
loaded = torch.load("price_mlp.pt", weights_only=True)
model = network(128)               # the shape comes from code
print(model.load_state_dict(loaded))
model.eval()                       # before the first prediction

checkpoint = {
    "epoch": epoch,
    "val_loss": val_loss,                 # a float, not np.float64
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),  # its state, by position
}
torch.save(checkpoint, "latest.pt")       # resuming only
if val_loss < best_loss:                  # ask before saving
    best_loss = val_loss
    torch.save(checkpoint, "best.pt")

back = torch.load("best.pt", weights_only=True)
model.load_state_dict(back["model"])
optimizer.load_state_dict(back["optimizer"])