Lesson 14

Checkpoints, inference, and responsible model reuse

Introduction

Four small classifier predictions contain three correct answers and one error, so accuracy is 0.75. Inspecting the wrong example matters before saving, loading, or fine-tuning the model for another task.

Learning goal

Evaluate, save, restore, and carefully fine-tune a PyTorch model while reporting prediction errors and important limitations.

Before you start

PyTorch modules, data splits, logits, class predictions, evaluation mode, and optimization.

Lesson plan

  1. Evaluate a classifier with aggregate accuracy and inspect individual prediction errors.
  2. Save configuration and learned state, reload them, and confirm matching outputs.
  3. Fine-tune selected parameters carefully and state data, metric, and reuse risks.

Why model use is more than one prediction

A trained model is only one part of a working system. Input preparation, tokenizer version, label meanings, evaluation data, and saved settings all affect the result. A model can run without an error while still producing wrong answers because its inputs were normalized differently or its class IDs were reversed.

Start from a model card or project documentation. Record what data the model expects, what task it was trained for, its limits, and its license. For text models, use the exact tokenizer and vocabulary that belong to the checkpoint. For image models, match resizing and normalization rules.

Practical terms

InferenceUsing learned parameters to produce predictions without training.
CheckpointSaved model state and, when needed, optimizer and training state.
state_dictA dictionary mapping parameter names to tensors.
Fine-tuningContinuing training from pretrained parameters on a new task or dataset.
Validation setHeld-out data used to compare choices during development.
Test setHeld-out data reserved for a final estimate.

Evaluate a small classifier

model.eval() switches layers such as dropout and batch normalization to evaluation behavior. It does not disable gradient recording. torch.no_grad() disables that recording for the enclosed operations and saves memory. Use both for ordinary evaluation.

Predict first: each fixed weight row selects one input feature. For [3,1], the logits are [3,1], so class 0 wins. Trace the other two rows before reading the printed predictions.

import torch
from torch import nn

model = nn.Linear(2, 2)
with torch.no_grad():
    model.weight.copy_(torch.tensor([[1.0, 0.0], [0.0, 1.0]]))
    model.bias.zero_()

features = torch.tensor([[3.0, 1.0],
                         [0.0, 2.0],
                         [4.0, 1.0]])
labels = torch.tensor([0, 1, 0])

model.eval()
with torch.no_grad():
    logits = model(features)
    predictions = logits.argmax(dim=1)
    accuracy = (predictions == labels).float().mean()

print(predictions.tolist())
print(accuracy.item())

# Expected:
# [0, 1, 0]
# 1.0

Accuracy is useful when classes are balanced and every error has similar cost. It can hide poor results on a rare class. Depending on the task, also inspect a confusion matrix, precision, recall, F1, calibration, or examples of failures. Measure against a simple baseline so that “high” has context.

Save and restore state

Saving a state_dict is usually clearer and more portable than saving a whole Python model object. To resume training exactly, save model state, optimizer state, the current epoch or step, and important configuration. Recreate the same architecture before loading its parameters.

A successful load message is not enough. Run the same fixed input through the original and restored models in evaluation mode. Equal outputs check the architecture and parameter transfer together.

import tempfile

optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

with tempfile.TemporaryDirectory() as directory:
    path = directory + "/checkpoint.pt"
    torch.save({
        "model": model.state_dict(),
        "optimizer": optimizer.state_dict(),
        "epoch": 3,
    }, path)

    restored_model = nn.Linear(2, 2)
    restored_optimizer = torch.optim.SGD(
        restored_model.parameters(), lr=0.01
    )
    checkpoint = torch.load(path, weights_only=True)
    restored_model.load_state_dict(checkpoint["model"])
    restored_optimizer.load_state_dict(checkpoint["optimizer"])
    restored_model.eval()

    with torch.no_grad():
        same = torch.equal(model(features), restored_model(features))
    print(checkpoint["epoch"])  # Expected: 3
    print(same)                 # Expected: True

A checkpoint does not automatically save the tokenizer, label mapping, random seeds, package versions, or model constructor arguments. Store these beside the checkpoint in a small configuration file or a documented experiment directory. Only load checkpoint files from sources you trust.

Fine-tune carefully

Fine-tuning starts from useful pretrained parameters instead of random values. Use a smaller learning rate than you might use from scratch, keep a validation split, and compare with a frozen-feature baseline. You can freeze a layer by setting its parameters' requires_grad to false, then give the optimizer only trainable parameters.

feature_layer = nn.Linear(4, 8)
classifier = nn.Linear(8, 2)

for parameter in feature_layer.parameters():
    parameter.requires_grad = False

trainable = [
    parameter
    for parameter in list(feature_layer.parameters())
                   + list(classifier.parameters())
    if parameter.requires_grad
]
print(len(trainable))  # Expected: 2 (classifier weight and bias)

Common pitfalls

  • model.eval() does not disable gradients; use no-grad or inference mode too.
  • Do not evaluate repeatedly on the final test set while choosing settings.
  • Match the checkpoint architecture, tokenizer, preprocessing, and label map.
  • Save optimizer state if training must resume, not only model weights.
  • Inspect performance across important groups and failure types, not only one average score.

Try it

A three-class model predicts [0,2,2,1] for labels [0,1,2,1]. Calculate accuracy. Which example is wrong, and why is looking at it useful?

Reveal the worked answer 
predictions = torch.tensor([0, 2, 2, 1])
labels = torch.tensor([0, 1, 2, 1])
correct = predictions == labels
print(correct.tolist())          # [True, False, True, True]
print(correct.float().mean().item())  # Expected: 0.75

The second example is wrong: class 2 was predicted instead of class 1. Examining errors can reveal ambiguous data, a label problem, missing training examples, or a systematic model weakness that one number cannot explain.

Recap

Reliable model use requires matching preprocessing, explicit evaluation, and reproducible saving. Use evaluation mode plus disabled gradients for inference. Save state dictionaries and enough configuration to rebuild the run. Fine-tune with held-out validation data and inspect failures as well as aggregate metrics. A model is ready only when its behavior is understood for the intended task.

References: PyTorch save/load tutorial and PyTorch no_grad documentation.