KAHIBARO
13.2 Loading for Inference and Resuming Training
Table of Contents
Two Different Goals, Two Different Loads
Loading a model usually means one of two things, and the steps differ depending on the goal. Inference loading is about making predictions reliably and efficiently. Resuming training is about continuing exactly where you left off, including optimizer state and training progress. Treat these as separate workflows so you do not accidentally load too little, or load the right tensors but in the wrong mode.
Loading for Inference
For inference, you typically only need the model architecture code and the trained weights. You reconstruct the model object in code, load its state_dict, move it to the correct device, switch it to evaluation mode, and run the forward pass without tracking gradients.
A standard pattern looks like this:
import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = MyModel(...) # same architecture arguments as training
state = torch.load("model_weights.pt", map_location=device)
model.load_state_dict(state)
model.to(device)
model.eval()
x = ... # torch.Tensor on the same device
x = x.to(device)
with torch.no_grad():
y = model(x)
Evaluation mode is important because certain layers behave differently at training time and inference time. If you forget to call model.eval(), you may see unstable predictions that do not match what you observed during validation.
Always call model.eval() for inference and wrap inference in with torch.no_grad(): unless you explicitly need gradients.
If you are loading on a machine without a GPU, using map_location="cpu" prevents errors and avoids accidental GPU allocations:
state = torch.load("model_weights.pt", map_location="cpu")When you need a numeric prediction as a Python value, detach it from PyTorch and move it to CPU first:
pred = y.detach().cpu()If your inference runs one sample at a time, remember to include a batch dimension if your model expects one. This is a common cause of shape errors.
If training used batched inputs of shape $(N, \dots)$, then single example inference usually needs shape $(1, \dots)$.
Loading to Resume Training
Resuming training is more than restoring weights. You generally want to restore at least the model parameters and the optimizer state. Often you also restore the epoch number, the best validation metric so far, a learning rate scheduler state, and any random number generator state if you care about continuing deterministically.
A typical checkpoint is a dictionary:
# during saving (covered in a different chapter)
checkpoint = {
"model_state": model.state_dict(),
"optimizer_state": optimizer.state_dict(),
"epoch": epoch,
"global_step": global_step,
"best_val": best_val,
}
torch.save(checkpoint, "checkpoint.pt")To resume:
import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = MyModel(...)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
ckpt = torch.load("checkpoint.pt", map_location=device)
model.load_state_dict(ckpt["model_state"])
optimizer.load_state_dict(ckpt["optimizer_state"])
model.to(device)
start_epoch = ckpt.get("epoch", 0) + 1
global_step = ckpt.get("global_step", 0)
best_val = ckpt.get("best_val", None)
model.train()
Calling model.train() matters here because you want training time behavior. Also note that the optimizer state contains tensors, for example running averages in Adam. After loading, ensure the optimizer state tensors are on the same device as the model parameters. Most of the time, moving the model to the device before creating the optimizer avoids confusion, but if you load the optimizer state and later change devices, you can end up with state tensors on the wrong device.
A safer ordering is to put the model on the target device before building the optimizer, then load states:
model = MyModel(...).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
ckpt = torch.load("checkpoint.pt", map_location=device)
model.load_state_dict(ckpt["model_state"])
optimizer.load_state_dict(ckpt["optimizer_state"])
To resume training correctly, you usually need both model.state_dict() and optimizer.state_dict(). Loading only model weights changes the training dynamics, especially for adaptive optimizers like Adam.
Resuming with a Learning Rate Scheduler
If you used a learning rate scheduler, restore it too. Otherwise, the learning rate schedule may restart or jump to an incorrect value.
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)
ckpt = torch.load("checkpoint.pt", map_location=device)
model.load_state_dict(ckpt["model_state"])
optimizer.load_state_dict(ckpt["optimizer_state"])
scheduler.load_state_dict(ckpt["scheduler_state"])If your scheduler depends on validation metrics, for example ReduceLROnPlateau, ensure you also restore any variables you use to track best metrics and when you last improved.
Common Loading Errors and How to Fix Them
A frequent error is missing keys or unexpected keys when calling load_state_dict. This usually means the model definition changed since the checkpoint was created, or you are loading weights into the wrong architecture. The correct fix is to ensure the architecture code and hyperparameters match the checkpoint’s original model.
Another common issue is saving with one device and loading on another. Using map_location prevents most device related loading errors.
A third issue appears when you wrap a model in torch.nn.DataParallel or similar wrappers during training, which can change the names of keys in the saved state_dict. If you trained with wrappers but infer without them, you may need to load the underlying module’s state_dict or adjust key names. The cleanest beginner friendly approach is to keep the saving and loading setup consistent between training and inference.
If load_state_dict reports many missing or unexpected keys, do not ignore it. It almost always means you are loading into a different model than the one that produced the checkpoint.
Minimal Checklist
For inference, reconstruct the model, load weights, move to device, call eval(), use torch.no_grad().
For resuming training, reconstruct model and optimizer, load both states, restore epoch or step counters, ensure everything is on the correct device, then call train() before continuing your training loop.
Views: 91