Skip to main content...
ML → Deep Learning via PyTorch — the Garment Classifier
20 min

Day 31: Checkpointing & inference mode

Saving, loading, and serving a trained model

A trained model is worthless if it evaporates when the process exits. Checkpointing saves the model's learned parameters (the state_dict) to disk so you can reload it for more training or for inference. The convention is to save the state_dict (just the weights) rather than the whole model object — it's portable and version-robust.

Saving and loading — state_dict is the portable way
# save
torch.save(model.state_dict(), "garment_classifier.pt")

# load for inference
model = GarmentClassifier(num_classes=8)   # recreate the architecture
model.load_state_dict(torch.load("garment_classifier.pt"))
model.eval()                               # inference mode

# a full checkpoint also saves optimizer state for resuming training
torch.save({
    "epoch": epoch,
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
}, "checkpoint.pt")

Two kinds of checkpoint, two purposes

For *deployment*, save just the model weights. For *resuming training*, also save the optimizer state (Adam's momentum buffers) and the epoch — otherwise resuming restarts the optimizer cold and the loss jumps. Knowing the difference signals you've actually run long training jobs, not just tutorials.

Key terms

state_dict
A dictionary mapping each layer to its learned parameter tensors — PyTorch's portable representation of a model's weights.
Checkpoint
A saved snapshot of training state (weights, optimizer state, epoch) allowing inference or resumed training.
Inference mode
Running a trained model to make predictions, with eval() set and gradients disabled.

You want to pause a long training run and resume it later exactly where it left off. Saving only model.state_dict() is insufficient. What else must you save?

We use cookies

We use cookies to enhance your browsing experience, serve personalized content, and analyze our traffic. By clicking "Accept All", you consent to our use of cookies. Learn more

    Day 31: Checkpointing & inference mode | RBTechIconX