PyTorch / Storage, restoration, and transformed execution / Lesson 14

State dictionaries and reviewed loading

Restore a defined function with compatible named state and an explicit loading policy.

About 4 min readContent reviewed 2026-10-28

On this page

Lesson goals

Distinguish a state_dict from architecture, preprocessing, and training continuation.
Preserve parameters and persistent buffers in a deliberate snapshot.
Use explicit load options and verify restored behavior.

A state dictionary has a specific scope

A module state_dict maps registered parameters and persistent buffers to named values. It does not by itself define the Python architecture, all preprocessing meaning, label mapping, or the complete state needed to resume training. Construct a compatible module and preserve the input/output contract alongside its state. A familiar filename and matching tensor dimensions cannot establish that columns retain the same meaning. For a training checkpoint, optimizer and scheduler state, random streams, data position, and partial accumulation may also matter. Choose contents from the intended restoration boundary rather than assuming one artifact serves every purpose.

A saved reference can change unless it is snapshotted

The illustrated program deep-copies the module state before a same-process in-memory save/load round trip. A state mapping contains tensor references, so merely keeping that mapping while continuing training is not automatically an immutable best-model snapshot. The registered center buffer is required with weight and bias to reconstruct this function. At input [3,2], subtracting [1,1] gives [2,1], whose weighted sum is three; adding bias 0.5 gives 3.5. No external checkpoint, filesystem model, or downloaded artifact is involved.

python Illustrative CPU PyTorch 2.14 API program; torch unavailable, expected output derived

import io, copy
import torch
from torch import nn
class Centered(nn.Module):
    def __init__(self):
        super().__init__()
        self.weight = nn.Parameter(torch.tensor([[2.0,-1.0]],dtype=torch.float64,device="cpu"))
        self.bias = nn.Parameter(torch.tensor([0.5],dtype=torch.float64,device="cpu"))
        self.register_buffer("center",torch.tensor([1.0,1.0],dtype=torch.float64,device="cpu"))
    def forward(self,x):
        return (x-self.center) @ self.weight.t()+self.bias
model = Centered()
snapshot = copy.deepcopy(model.state_dict())
stream = io.BytesIO()
torch.save(snapshot,stream)
stream.seek(0)
state = torch.load(stream,map_location="cpu",weights_only=True)
restored = Centered()
restored.load_state_dict(state,strict=True)
restored.eval()
x = torch.tensor([[3.0,2.0]],dtype=torch.float64,device="cpu")
with torch.no_grad():
    output = restored(x)
print("keys", sorted(state))
print("output", output.tolist())

Expected output

keys ['bias', 'center', 'weight']
output [[3.5]]

Only a self-created in-memory tensor state is described. Explicit weights_only and CPU mapping do not imply general safe loading or compatibility; no save/load or model call was executed.

Loading options do not replace a trust policy

The fixture explicitly requests map_location="cpu" and weights_only=True rather than relying on an evolving default. The documented restricted loading behavior reduces what the loader accepts; it is not a complete guarantee against every malicious or oversized input, nor proof of trusted origin. Do not disable restrictions or allow unknown globals merely to make an unfamiliar file load. Review artifact provenance, format, expected keys, types, shapes, resource bounds, and supported dependencies. This course does not open unknown serialized objects and makes no security certification or runtime compatibility claim.

Strict key matching is one check among several

load_state_dict with strict=True checks the expected registered key relationship under its documented contract. It does not prove feature units, class order, or evaluation validity. A deliberate partial load needs an explicit missing/unexpected-key policy and separate assessment of resulting behavior. The standard-library companion exposes a missing center value, restores a self-created JSON mapping, and checks the resulting scalar arithmetic. It does not reproduce tensor metadata, storage sharing, checkpoint versions, optimizer matching, or torch loading. Keep structural checks and known-input behavior checks distinct and both visible.

python Executed deterministic standard-library companion

import json
required = {"weight","bias","center"}
incomplete = {"weight":[2.0,-1.0],"bias":0.5}
print("missing state", sorted(required-incomplete.keys()))
complete = dict(incomplete,center=[1.0,1.0])
restored = json.loads(json.dumps(complete,allow_nan=False))
x = [3.0,2.0]
value = sum((a-c)*w for a,c,w in zip(x,restored["center"],restored["weight"]))+restored["bias"]
print(f"restored scalar={value:.2f}")

Expected output

missing state ['center']
restored scalar=3.50

The executed JSON fixture checks required named state and the declared scalar function. It is not a torch checkpoint or a serialization compatibility test.

Verify the restored inference and continuation boundary

Set intended module mode and gradient context, then compare known raw inputs through the full preprocessing-to-output path with the reference. Include threshold-near and unsupported inputs where decisions require them. If resuming training, compare the next input, gradient, optimizer transition, and schedule instead of only one inference score. Record exactly which environment and files were tested. The PyTorch program remains illustrative because torch is unavailable; only JSON and scalar arithmetic were executed. No state_dict round trip, tensor loader, framework checkpoint, model inference, or training continuation was run during review.

Common errors and fixes

The remembered best state changes during later training

Why it happens: The state mapping retained references rather than an independent snapshot.
How to fix it: Copy or save at the intended boundary with a reviewed policy.
Prevent it next time: Test a controlled model change after snapshot creation.

A checkpoint loads but predictions have another meaning

Why it happens: Architecture or input/output semantics changed despite compatible keys.
How to fix it: Validate the complete model and preprocessing contract.
Prevent it next time: Keep named-state checks and a known raw-input fixture.

A restricted loader error is bypassed without review

Why it happens: An unsupported artifact was treated as a reason to remove loading safeguards.
How to fix it: Establish trusted provenance and the documented supported format first.
Prevent it next time: Test missing, unexpected, incompatible, and oversized state under explicit policies.

Check your understanding

1. Why must center be saved even though it is not an optimizer parameter?

Show worked solution

The forward subtracts that persistent buffer before applying the weight. Without its value, the same input and trainable parameters can produce another function. Persistence and optimization are different roles, so named_parameters alone is not the complete inference state.

2. Would strict=True and weights_only=True prove the package is appropriate for a task?