Table of Contents
What to Save and Why
In PyTorch you typically save either model weights only or a full training checkpoint. Saving weights only means saving the parameters of your neural network, usually with model.state_dict(). This is the most common choice for deployment and inference because it is compact and portable across training setups. A full checkpoint saves enough information to resume training later as if nothing happened. That usually includes the model weights, optimizer state, learning rate scheduler state, the current epoch or step, and any metadata you need to reproduce the run.
Rule of thumb: save weights for inference, save full checkpoints for resuming training.
Saving Model Weights with `state_dict`
A model’s state_dict is a Python dictionary that maps parameter and buffer names to tensors. You save it with torch.save and later load it into an identically constructed model.
import torch
# during or after training
torch.save(model.state_dict(), "model_weights.pt")When loading, you must first create the model object with the same architecture, then load the weights.
model = MyModel(**same_constructor_args)
state = torch.load("model_weights.pt", map_location="cpu")
model.load_state_dict(state)
model.eval()
Calling model.eval() matters for inference because layers like dropout and batch normalization behave differently in training mode versus evaluation mode.
Always reconstruct the model architecture before calling load_state_dict. The saved state_dict does not contain the Python code that defines your model.
Saving Full Checkpoints for Resuming Training
A full checkpoint is usually a dictionary you define yourself. At minimum it includes the model and optimizer states. If you use a learning rate scheduler, include that too. It is also common to store the epoch or global step and the best validation metric so far.
checkpoint = {
"epoch": epoch,
"model_state": model.state_dict(),
"optimizer_state": optimizer.state_dict(),
"scheduler_state": scheduler.state_dict() if scheduler is not None else None,
"best_val_loss": best_val_loss,
}
torch.save(checkpoint, "checkpoint.pt")To resume, you load the dictionary, restore states, and continue from the saved epoch.
ckpt = torch.load("checkpoint.pt", map_location=device)
model.load_state_dict(ckpt["model_state"])
optimizer.load_state_dict(ckpt["optimizer_state"])
if scheduler is not None and ckpt["scheduler_state"] is not None:
scheduler.load_state_dict(ckpt["scheduler_state"])
start_epoch = ckpt["epoch"] + 1
best_val_loss = ckpt.get("best_val_loss", float("inf"))If you want to resume training faithfully, you must save and load the optimizer state. Without it, momentum based optimizers like Adam will behave as if training just started, even if the weights look correct.
Device and `map_location` Details
Checkpoints can be saved on one device and loaded on another. Use map_location to control where tensors land when loading. A common safe pattern is to load on CPU first, then move the model to the target device.
ckpt = torch.load("checkpoint.pt", map_location="cpu")
model.load_state_dict(ckpt["model_state"])
model.to(device)This reduces the chance of GPU memory spikes during loading.
Naming Conventions and What to Keep
It helps to maintain two kinds of files: a latest checkpoint that you overwrite each epoch, and a best checkpoint that you only update when validation improves. The latest file is for recovery from interruptions, the best file is for evaluation and deployment.
Do not rely on only one checkpoint file. Keep at least a latest and a best to protect against crashes and bad epochs.
Common Pitfalls
One frequent issue is shape mismatches when loading weights, which usually means the model definition changed. Another is forgetting to save non model state needed to resume training, such as the scheduler. Also note that some layers register non parameter buffers, for example batch normalization running statistics, and these are included in state_dict, which is one reason state_dict is preferred over saving individual parameter tensors.
When something does not load, inspect the keys.
state = torch.load("model_weights.pt", map_location="cpu")
print(list(state.keys())[:10])Recommended Minimal Templates
For inference, the minimal template is saving model.state_dict() and separately recording the model constructor arguments or configuration in a JSON or Python dict. For training resumption, the minimal template is a dict with model_state, optimizer_state, and an epoch or step counter, plus scheduler state if you use one.
A checkpoint is only as useful as the metadata you store with it. If you cannot reconstruct the model configuration, the weights are not reusable.