🎓 A1Learner
🔍

PyTorch Saving and Loading Models

Saving a model means saving its state_dict, the table of its learned numbers. Loading copies those numbers into a fresh model of the same shape.

Save the numbers

One call writes the state_dict to a file. The file lands in the working folder.

torch.save(model.state_dict(), "model.pt")

Load into a fresh model

The fresh model starts with different random numbers. Loading overwrites them with the saved ones.

fresh = nn.Linear(1, 1)
fresh.load_state_dict(torch.load("model.pt"))
print("Weights equal:", torch.equal(model.weight, fresh.weight))
print("Bias equal:", torch.equal(model.bias, fresh.bias))

Example

import torch
from torch import nn

torch.manual_seed(0)
x = torch.tensor([[1.0], [2.0], [3.0], [4.0]])
y = torch.tensor([[3.0], [5.0], [7.0], [9.0]])

torch.manual_seed(0)
model = nn.Linear(1, 1)
loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.05)

for step in range(200):
    pred = model(x)
    loss = loss_fn(pred, y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

torch.save(model.state_dict(), "regressor.pt")

torch.manual_seed(99)
fresh = nn.Linear(1, 1)
fresh.load_state_dict(torch.load("regressor.pt"))
fresh.eval()
model.eval()

test = torch.tensor([[5.0], [6.0]])
with torch.no_grad():
    first = model(test)
    second = fresh(test)

print("Original predictions:", [round(v, 2) for v in first.flatten().tolist()])
print("Loaded predictions:", [round(v, 2) for v in second.flatten().tolist()])
print("Identical:", torch.equal(first, second))

Key points

  • state_dict: the model's learned numbers held in one table
  • Use torch.save(model.state_dict(), path) to save the numbers
  • Use fresh.load_state_dict(torch.load(path)) to load them into a model
  • The fresh model must have the same layer shapes as the saved one
  • Saved numbers copy exactly, so predictions match exactly
  • A state_dict carries the trained result, ready for a new program

In this example, seed 0 is set before the data and again before the model. The model trains for 200 steps on the line y = 2x + 1. Its state_dict is saved to the file regressor.pt. Seed 99 then builds a fresh Linear layer. That seed gives the fresh layer different starting numbers. load_state_dict copies the saved numbers into it. Both models switch to eval mode. A no_grad pass runs two test inputs through each model. The two prediction lists print the same. torch.equal confirms the outputs are identical. The saved numbers carried the whole trained model.

Note

The state_dict saves numbers only, not the model's code

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What must be true of a fresh model before load_state_dict can fill it?