Saving & Loading Models

By the end of this lesson you will be able to save a trained PyTorch model and load it back to make predictions, write a full checkpoint that lets you resume an interrupted training run, and load files safely. You will also know what extra information to save so the model is still usable months later.

Training can take hours or days. If you cannot save the result, all of that is lost when the program ends. And a saved model is what you will deploy. The ideas follow Saving & Loading Models from the machine learning course, adapted for PyTorch.

What to save: the state dict

A PyTorch model has two parts:

1. the code that defines its layers (your nn.Module class);
2. the numbers learned in training: weights, biases and buffers such as batch norm running statistics.

The recommended way is to save only the numbers, as the model's state dict: an ordered dictionary from layer names to tensors. To load, you create the model from your code and pour the saved numbers back in.

Think of a family recipe for rajma. The method (soak, boil, fry the masala, simmer) lives in the recipe book, which is your code. The exact amounts your grandmother perfected over years (how much salt, how many minutes) are the numbers. You only need to write down the amounts, as long as you keep the recipe book.

Step 1: save and load a state dict. This is plain PyTorch code with its real output.

import torch
from torch import nn

def make_model():
    return nn.Sequential(nn.Linear(4, 16), nn.ReLU(), nn.Linear(16, 3))

torch.manual_seed(0)
model = make_model()
print("keys:", list(model.state_dict().keys()))
torch.save(model.state_dict(), "model_weights.pt")

loaded = make_model()                                  # same code, new random weights
loaded.load_state_dict(torch.load("model_weights.pt"))
loaded.eval()

x = torch.randn(2, 4)
print("same predictions:", torch.equal(model(x), loaded(x)))

keys: ['0.weight', '0.bias', '2.weight', '2.bias']
same predictions: True

The keys such as 0.weight come from the layer positions in the Sequential. If you change the model code so the names or shapes no longer match, load_state_dict raises an error listing the missing and unexpected keys. That strictness is helpful: it catches mismatches early.

Always call model.eval() after loading for prediction, so dropout and batch norm behave correctly.

Saving the whole model object

torch.save(model, "model.pt") also works. It pickles the whole object, including a reference to your class. It looks simpler, but it breaks easily: loading it later needs the exact same class, in a module with the same name, in the same place. Rename a file and the saved model can no longer be loaded. Prefer the state dict.

Loading safely

torch.save uses Python's pickle format. Unpickling a file can run arbitrary code, so a malicious model file could harm your computer. Since PyTorch 2.6, torch.load uses weights_only=True by default: it only accepts tensors, plain containers and a few other safe types, and refuses anything else with an error.

import torch

class Something:
    pass

torch.save({"weights": torch.ones(2), "extra": Something()}, "unusual.pt")
try:
    torch.load("unusual.pt")
except Exception as err:
    print(type(err).__name__, "- the file contains a custom object, so it was refused")

print(torch.load("model_weights.pt").keys())            # a plain state dict loads fine

UnpicklingError - the file contains a custom object, so it was refused
odict_keys(['0.weight', '0.bias', '2.weight', '2.bias'])

You can pass weights_only=False to load such a file, but only do that for files you created yourself or fully trust. Never do it for a model file downloaded from an unknown source. Many model hubs also offer the safetensors format, which stores only tensors and cannot run code at all.

Step 2: checkpoints to resume training

If training stops halfway (a power cut in Kochi, a cloud notebook timing out), you want to continue from where you were, not from scratch. For that you need more than the weights:

• the model's state dict;
• the optimiser's state dict (Adam's running averages, the current learning rate);
• the scheduler's state, if you use one;
• the epoch number and the best score so far.

Put them in one dictionary and save it. This is plain code with its real output.

import torch
from torch import nn

torch.manual_seed(0)
X, y = torch.randn(256, 4), torch.randint(0, 3, (256,))
model = nn.Sequential(nn.Linear(4, 16), nn.ReLU(), nn.Linear(16, 3))
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

def train_epochs(model, optimizer, start, end):
    for epoch in range(start, end):
        loss = nn.functional.cross_entropy(model(X), y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    return loss.item()

loss = train_epochs(model, optimizer, 0, 50)
torch.save({"epoch": 50, "model": model.state_dict(),
            "optimizer": optimizer.state_dict(), "loss": loss}, "checkpoint.pt")
print(f"saved at epoch 50, loss {loss:.4f}")

# ... later, perhaps in a new program ...
model2 = nn.Sequential(nn.Linear(4, 16), nn.ReLU(), nn.Linear(16, 3))
optimizer2 = torch.optim.Adam(model2.parameters(), lr=0.01)
ckpt = torch.load("checkpoint.pt")
model2.load_state_dict(ckpt["model"])
optimizer2.load_state_dict(ckpt["optimizer"])
model2.train()
loss2 = train_epochs(model2, optimizer2, ckpt["epoch"], 100)

loss1 = train_epochs(model, optimizer, 50, 100)              # the run that never stopped
print(f"resumed run at epoch 100: loss {loss2:.4f}")
print(f"unbroken run at epoch 100: loss {loss1:.4f}")

saved at epoch 50, loss 1.0276
resumed run at epoch 100: loss 0.9793
unbroken run at epoch 100: loss 0.9793

The resumed run reaches exactly the same loss as the run that never stopped, because the optimiser state came back too. If you restore only the model weights, Adam restarts with empty averages and the run takes a slightly different path.

Step 3: the same idea in NumPy

There is nothing special about the file format. Our NumPy network from A Two-Layer Network from Scratch is just a few arrays, and NumPy's savez stores named arrays in one file. You can run this here; the file lives in the browser's temporary folder.

import json
import numpy as np

rng = np.random.default_rng(0)
params = {"W1": rng.normal(0, 0.5, (64, 32)), "b1": np.zeros(32),
          "W2": rng.normal(0, 0.5, (32, 10)), "b2": np.zeros(10)}

def predict(p, X):
    return (np.maximum(0, X @ p["W1"] + p["b1"]) @ p["W2"] + p["b2"]).argmax(axis=1)

np.savez("digit_net.npz", **params)                        # weights
meta = {"classes": list(range(10)), "input_scale": 16.0, "version": 1}
with open("digit_net.json", "w", encoding="utf-8") as f:
    json.dump(meta, f)                                     # everything else you need

loaded = dict(np.load("digit_net.npz"))
with open("digit_net.json", encoding="utf-8") as f:
    meta2 = json.load(f)

X = rng.random((5, 64)) * 16 / meta2["input_scale"]
print("arrays:", {k: v.shape for k, v in loaded.items()})
print("same predictions:", np.array_equal(predict(params, X), predict(loaded, X)))
print("metadata:", meta2)

arrays: {'W1': (64, 32), 'b1': (32,), 'W2': (32, 10), 'b2': (10,)}
same predictions: True
metadata: {'classes': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9], 'input_scale': 16.0, 'version': 1}

Save the metadata too

A file of weights is useless if nobody remembers how to use it. Next to every model you plan to keep, save:

1. the class names in output order (index 3 means "Dress", not just "3");
2. the preprocessing: image size, normalisation mean and standard deviation, vocabulary;
3. the architecture settings used to build the model (layer sizes, number of classes);
4. the library versions (PyTorch, torchvision);
5. the validation and test scores and the date.

You can put these in the checkpoint dictionary or in a small JSON file beside it, as in Step 3.

Moving between CPU and GPU

A model trained on a GPU stores its tensors as GPU tensors. Loading that file on a computer without a GPU fails unless you tell torch.load where to put them:

state = torch.load("best_model.pt", map_location="cpu")    # GPU-trained file on a CPU machine
model.load_state_dict(state)
model.to(device)

Common mistakes

• Forgetting model.eval() after loading. Predictions vary with dropout on.
• Saving only the weights when you want to resume. Include the optimiser, scheduler and epoch.
• Loading untrusted files with weights_only=False. Pickle can run code; only load files you trust.
• Changing the model code after saving. Renamed layers no longer match the state dict keys.
• Not saving class names and preprocessing. The model gives numbers nobody can interpret.

Key points

• Save model.state_dict() and load it into a model built from the same code.
• A checkpoint for resuming holds the model, optimiser and scheduler states plus the epoch and best score.
• torch.load uses weights_only=True by default since PyTorch 2.6; never disable it for untrusted files.
• Use map_location to load GPU-trained weights on a CPU.
• Save metadata (classes, preprocessing, settings, versions, scores) next to the weights.

import torch
from torch import nn

torch.manual_seed(0)
X, y = torch.randn(256, 4), torch.randint(0, 3, (256,))
model = nn.Sequential(nn.Linear(4, 16), nn.ReLU(), nn.Linear(16, 3))
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.5)
for epoch in range(50):
    loss = nn.functional.cross_entropy(model(X), y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    scheduler.step()
torch.save({"epoch": 50, "model": model.state_dict(), "optimizer": optimizer.state_dict(),
            "scheduler": scheduler.state_dict()}, "checkpoint_lr.pt")

model2 = nn.Sequential(nn.Linear(4, 16), nn.ReLU(), nn.Linear(16, 3))
optimizer2 = torch.optim.Adam(model2.parameters(), lr=0.01)
scheduler2 = torch.optim.lr_scheduler.StepLR(optimizer2, step_size=30, gamma=0.5)
ckpt = torch.load("checkpoint_lr.pt")
model2.load_state_dict(ckpt["model"])
optimizer2.load_state_dict(ckpt["optimizer"])
scheduler2.load_state_dict(ckpt["scheduler"])
print("resumed learning rate:", optimizer2.param_groups[0]["lr"])
print("scheduler last lr:   ", scheduler2.get_last_lr())

resumed learning rate: 0.005
scheduler last lr:    [0.005]