PyTorch: state_dict checkpoints and explicit safe loading

pythontorchindependent-practice

Saved on this device

Full lesson · 4 learning steps Open any exercise for focused practice

Step 1 of 4 · intro card

PyTorch: state_dict checkpoints and explicit safe loading

Concept and contract

Restore owned weights into a matching architecture without loading arbitrary objects.

A model checkpoint should have a clear producer, architecture version and loading policy. state_dict captures registered parameter and buffer values rather than requiring serialization of a whole application object.

Reconstruct the intended architecture explicitly, load its state and choose device placement deliberately. A saved file is still an input with provenance and resource risks, not automatically trustworthy because it has a familiar extension.

Mechanism and worked example

The helper writes a known one-weight model’s state into an owned BytesIO buffer, rewinds it and loads with weights_only=True and map_location CPU.

The helper writes a known one-weight model’s state into an owned BytesIO buffer, rewinds it and loads with weights_only=True and map_location CPU. A fresh matching Linear module receives the state and performs evaluation.

The returned numerical prediction demonstrates that restoration actually occurred. No downloaded checkpoint, arbitrary pickle object or production file is opened.

Complete worked example

python

from bootstrap_lab import torch_setup
torch = torch_setup()

import io
from torch import nn

def checkpoint_roundtrip(weight, value):
    source = nn.Linear(1, 1, bias=False)
    with torch.no_grad():
        source.weight.fill_(weight)
    buffer = io.BytesIO()
    torch.save(source.state_dict(), buffer)
    buffer.seek(0)
    restored = nn.Linear(1, 1, bias=False)
    state = torch.load(buffer, map_location='cpu', weights_only=True)
    restored.load_state_dict(state)
    restored.eval()
    with torch.inference_mode():
        return restored(torch.tensor([[float(value)]])).item()

Reasoning: The helper writes a known one-weight model’s state into an owned BytesIO buffer, rewinds it and loads with weights_only=True and map_location CPU. A fresh matching Linear module receives the state and performs evaluation.

Counterexamples and common mistakes

The counterexample loads the dictionary but never applies it to the new module, leaving random weights in use.

The counterexample loads the dictionary but never applies it to the new module, leaving random weights in use. A successful file read alone cannot prove restored behavior.

Avoid loading untrusted general Python objects with unrestricted pickle semantics. Explicit weights_only is a narrower loading policy, not a complete defense against every malformed file, resource exhaustion or dishonest model artifact.

Counterexample: a runnable framework mistake

python

from bootstrap_lab import torch_setup
torch = torch_setup()

import io
from torch import nn

def checkpoint_roundtrip(weight, value):
    source = nn.Linear(1, 1, bias=False)
    with torch.no_grad():
        source.weight.fill_(weight)
    buffer = io.BytesIO()
    torch.save(source.state_dict(), buffer)
    buffer.seek(0)
    restored = nn.Linear(1, 1, bias=False)
    state = torch.load(buffer, map_location='cpu', weights_only=True)
    pass  # state never applied
    restored.eval()
    with torch.inference_mode():
        return restored(torch.tensor([[float(value)]])).item()

Reasoning: The counterexample loads the dictionary but never applies it to the new module, leaving random weights in use. A successful file read alone cannot prove restored behavior.

Independent practice: one bounded framework change

Also return the restored weight value beside its prediction.

Also return the restored weight value beside its prediction. Preserve the owned buffer, matching architecture, explicit load options and evaluation context.

This small readback addition ties the observed prediction to the loaded parameter. Do not return the source argument as a shortcut: inspect the actual restored module so the fixture can expose missing state application.

Step 1: Return the actual restored weight beside its prediction. Step 2: Read restored.weight after load_state_dict.

Practice starter: solve before reading the solution

python

from bootstrap_lab import torch_setup
torch = torch_setup()

import io
from torch import nn

def checkpoint_roundtrip(weight, value):
    # Remove this placeholder, then adapt the retained baseline:
    # Return the actual restored weight beside its prediction.
    raise NotImplementedError
    source = nn.Linear(1, 1, bias=False)
    with torch.no_grad():
        source.weight.fill_(weight)
    buffer = io.BytesIO()
    torch.save(source.state_dict(), buffer)
    buffer.seek(0)
    restored = nn.Linear(1, 1, bias=False)
    state = torch.load(buffer, map_location='cpu', weights_only=True)
    restored.load_state_dict(state)
    restored.eval()
    with torch.inference_mode():
        return restored(torch.tensor([[float(value)]])).item()

Reasoning: Return the actual restored weight beside its prediction. Read restored.weight after load_state_dict.

Hint 1: Keep the explicit CPU and weights-only options.
Hint 2: Do not echo the input argument as proof of restoration.

Solution, tests and trade-offs

Cases use positive, negative and zero weights and compare resulting predictions.

Cases use positive, negative and zero weights and compare resulting predictions. Practice also inspects the restored weight.

Finite tolerances account for floating representation. The tests establish a small owned local round trip only; they do not certify arbitrary third-party checkpoints, architecture migrations, optimizer-state resumption or compatibility across every future PyTorch release.

Practice solution and acceptance checks

python

from bootstrap_lab import torch_setup
torch = torch_setup()

import io
from torch import nn

def checkpoint_roundtrip(weight, value):
    source = nn.Linear(1, 1, bias=False)
    with torch.no_grad():
        source.weight.fill_(weight)
    buffer = io.BytesIO()
    torch.save(source.state_dict(), buffer)
    buffer.seek(0)
    restored = nn.Linear(1, 1, bias=False)
    state = torch.load(buffer, map_location='cpu', weights_only=True)
    restored.load_state_dict(state)
    restored.eval()
    with torch.inference_mode():
        return {'prediction': restored(torch.tensor([[float(value)]])).item(), 'weight': restored.weight.item()}

Reasoning: Cases use positive, negative and zero weights and compare resulting predictions. Practice also inspects the restored weight.

Interview reasoning and next step

For resumable training, save the required optimizer state, epoch and configuration with a documented schema and trusted provenance.

For resumable training, save the required optimizer state, epoch and configuration with a documented schema and trusted provenance. Keep inference-only exports distinct from training checkpoints.

If architecture changes produce missing or unexpected keys, review them intentionally rather than suppressing strict loading without explanation. Preserve preprocessing parameters and evaluation mode as part of the complete inference pipeline.

Optional stretch — after the core task works: Define a versioned training checkpoint with owned optimizer and normalization state.

Computer practice: https://applydjinn.com/typorithm/python-frameworks-lab-20261105.zip. Install the pinned dependencies in an isolated environment, run verify.py --track torch, then edit the retained exercise baseline and check it with --practice pt-checkpoints. The lab uses owned in-memory databases, local framework clients and tiny CPU tensors. It never reads the working ApplyDjinn database, contacts production feeds or downloads models.

Restore owned weights into a matching architecture without loading arbitrary objects.