Save, Load, Export, and Version Models

Save, Load, Export, and Version Models is the point where a PyTorch experiment becomes a reproducible artifact. In this lesson, the outcome is concrete: you will know what to store, how PyTorch reconstructs model state, when to export an inference graph, and how to attach enough version information that a deployed model can be compared, replaced, or rolled back.

This chapter fits the deployment part of the Deep Learning with PyTorch course. Earlier lessons focused on tensors, modules, loss functions, optimizers, and evaluation. Those pieces only survive beyond a notebook if you serialize the right state and reload it into compatible code. A saved file is not just a convenience; it is the bridge between training, evaluation, batch scoring, serving, and incident response.

Purpose and Outcome

Saving a model answers, “Can I use these learned parameters again?” Loading answers, “Can I reconstruct the same computation in another process?” Export answers, “Can I run inference without depending on the exact training program?” Versioning answers, “Can I prove which artifact produced this prediction and recover the previous one?” A complete workflow normally includes all four.

The practical goal is to separate architecture, parameters, training state, and serving artifact. Architecture is Python code such as an nn.Module class. Parameters and buffers are tensors held in a module’s state_dict. Training state includes optimizer momentum, scheduler position, epoch number, random number generator state, and metrics. Serving artifacts may be a PyTorch state dictionary loaded by application code, a TorchScript module, or another export format chosen by the serving platform.

How PyTorch Serialization Works

Every nn.Module recursively owns parameters, such as layer weights, and buffers, such as batch-normalization running means. Calling model.state_dict() returns an ordered mapping from stable string names to tensors. A key like encoder.layers.0.self_attn.in_proj_weight records where the tensor lives in the module tree. torch.save serializes that mapping to a file, and load_state_dict copies tensors from the mapping into an already constructed module.

That last detail matters. A state dictionary does not contain the Python class definition or the constructor arguments that built the model. You recreate the architecture first, then load tensors into matching names and shapes. If names are missing, unexpected, or shape-incompatible, PyTorch reports the mismatch. This design keeps weight files smaller and more portable across refactors than pickling a whole Python object, but it also means you must preserve or intentionally migrate the module naming scheme.

torch.save and torch.load can store arbitrary Python objects, including a whole module, but whole-object pickle files bind the artifact to import paths and executable Python behavior. For long-lived lessons, teams, and deployments, the usual PyTorch recommendation is to save state dictionaries and a small, explicit manifest rather than relying on pickled model objects.

API Anatomy

The core API has a small surface area. torch.save(obj, path) writes a serialized object. torch.load(path, map_location="cpu") reads it while remapping tensor storage to the requested device. model.state_dict() extracts named tensors. model.load_state_dict(state, strict=True) validates that the keys match exactly by default. For checkpoints, you usually save a dictionary with several named entries rather than saving only the model weights.

Device handling is a common source of confusion. A tensor saved from CUDA may try to load onto CUDA unless you use map_location. Load to CPU first when portability matters, then move the model to the inference or training device with model.to(device). Also switch modes deliberately: model.train() enables training behavior such as dropout, while model.eval() disables dropout and uses stored normalization statistics for inference.

Example 1: Save and Reload Weights

The first example saves only learned module state. It creates the same class twice: one instance is trained in principle, and the second receives the saved tensors. The deterministic checks are the missing and unexpected key lists, followed by whether both models produce the same output for the same input.

import tempfile
from pathlib import Path

import torch
from torch import nn

class TinyClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(4, 3), nn.ReLU(), nn.Linear(3, 2))

    def forward(self, x):
        return self.net(x)

torch.manual_seed(7)
model = TinyClassifier()
example = torch.randn(1, 4)
with tempfile.TemporaryDirectory() as tmp:
    path = Path(tmp) / "tiny-state.pt"
    torch.save(model.state_dict(), path)

    restored = TinyClassifier()
    incompatible = restored.load_state_dict(torch.load(path, map_location="cpu"))
    same_output = torch.allclose(model(example), restored(example))
    print(incompatible.missing_keys, incompatible.unexpected_keys)
    print(same_output)

Expected output is [] [] and then True. The empty lists mean the receiving architecture had every key required by the file and the file had no extra keys. The output comparison succeeds because the two module instances now hold identical tensor values. In a real project, use this pattern for inference services that can import the model class and only need weights.

Example 2: Checkpoint Training State

Saving weights alone is not enough to resume training faithfully. Optimizers keep internal state, such as momentum buffers or adaptive learning-rate statistics. Schedulers, epoch counters, scaler state for mixed precision, and random number generator state can also change the next update. A checkpoint is a structured dictionary that records the training process, not just the network.

import tempfile
from pathlib import Path

import torch
from torch import nn

model = nn.Linear(3, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
x = torch.ones(5, 3)
y = torch.ones(5, 1)
loss = nn.functional.mse_loss(model(x), y)
loss.backward()
optimizer.step()

checkpoint = {
    "epoch": 4,
    "model_state": model.state_dict(),
    "optimizer_state": optimizer.state_dict(),
    "rng_state": torch.random.get_rng_state(),
    "metrics": {"val_loss": 0.37},
}
with tempfile.TemporaryDirectory() as tmp:
    path = Path(tmp) / "checkpoint.pt"
    torch.save(checkpoint, path)
    loaded = torch.load(path, map_location="cpu")

resumed = nn.Linear(3, 1)
resumed.load_state_dict(loaded["model_state"])
resumed_optimizer = torch.optim.SGD(resumed.parameters(), lr=0.1, momentum=0.9)
resumed_optimizer.load_state_dict(loaded["optimizer_state"])
torch.random.set_rng_state(loaded["rng_state"])
print(loaded["epoch"])
print(sorted(loaded["metrics"]))

The script prints 4 and then ['val_loss']. The resumed model receives the tensor weights, and the resumed optimizer receives its momentum state. In practice, continue from epoch + 1, restore any scheduler or gradient-scaler state you use, and record the validation metric that justified keeping the checkpoint. This example is the minimum shape of a restartable training run.

Example 3: Export an Inference Artifact

A deployment target may not want the whole training program. TorchScript captures a callable module representation that can be saved and loaded independently of the original eager Python object. Tracing records operations observed for example inputs, so it is best for models whose control flow does not depend on input values. Scripting analyzes Python control flow more directly, but it requires code that TorchScript can compile.

import tempfile
from pathlib import Path

import torch
from torch import nn

class ScoreModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(2, 1)

    def forward(self, features):
        return torch.sigmoid(self.linear(features))

torch.manual_seed(3)
model = ScoreModel().eval()
example = torch.zeros(1, 2)
with torch.no_grad():
    traced = torch.jit.trace(model, example)

with tempfile.TemporaryDirectory() as tmp:
    path = Path(tmp) / "score-model.pt"
    traced.save(str(path))
    loaded = torch.jit.load(str(path), map_location="cpu")
    eager_value = model(example)
    exported_value = loaded(example)
    print(torch.allclose(eager_value, exported_value))

The expected output is True. The eager model and loaded TorchScript module produce the same value for the example input. This does not prove the export is correct for every possible input; it proves the artifact round-trips and preserves the traced computation for the checked case. For production, test representative shapes, boundary values, and the exact preprocessing path used by clients.

Example 4: Version with a Manifest

Versioning should include more than a filename. A manifest is a small machine-readable document next to the artifact. It identifies the artifact, expected model class, input schema, validation metrics, and a cryptographic digest of the weight file. The digest detects accidental replacement or corruption.

import hashlib
import json
import tempfile
from pathlib import Path

import torch
from torch import nn

model = nn.Linear(2, 2)
with tempfile.TemporaryDirectory() as tmp:
    root = Path(tmp)
    weights_path = root / "model_state.pt"
    torch.save(model.state_dict(), weights_path)
    digest = hashlib.sha256(weights_path.read_bytes()).hexdigest()
    manifest = {
        "artifact_id": "fraud-score-demo",
        "artifact_format": "pytorch-state-dict",
        "model_class": "torch.nn.Linear",
        "input_schema": {"features": ["amount_zscore", "merchant_risk"]},
        "metrics": {"validation_auc": 0.91},
        "weights_sha256": digest,
    }
    manifest_path = root / "manifest.json"
    manifest_path.write_text(json.dumps(manifest, sort_keys=True), encoding="utf-8")
    loaded = json.loads(manifest_path.read_text(encoding="utf-8"))
    print(loaded["artifact_id"])
    print(loaded["weights_sha256"] == digest)

The expected output is fraud-score-demo and then True. The artifact can now be checked before loading: compute the current file digest, compare it with weights_sha256, verify the input schema, and confirm the metric or approval record. This pattern scales from local folders to object storage and model registries.

Design Choices and Trade-offs

State dictionary versus whole module: state dictionaries require the code to define the architecture, but they are easier to inspect, migrate, and review. Whole modules are convenient for quick experiments, but they are brittle across import-path changes and risky because loading pickle data can execute code.

Checkpoint versus inference artifact: checkpoints preserve training context and may include optimizer state that is irrelevant to serving. Inference artifacts should be smaller, evaluated in eval mode, and paired with preprocessing and postprocessing definitions. Keeping both is normal: one lets you continue training, the other lets you serve predictions.

Strict loading versus partial loading: strict=True catches accidental architecture drift. strict=False can be useful for transfer learning when replacing a classifier head, but every missing or unexpected key must be reviewed. Treat partial loading as an explicit migration, not as a way to silence errors.

TorchScript or eager PyTorch: eager loading is flexible and easy to debug inside Python services. TorchScript can reduce dependency on Python source at serving time and can be used by runtimes that expect a serialized module. Tracing is simple but may miss data-dependent branches; scripting is stronger for control flow but more restrictive.

Failure Modes and Troubleshooting

Symptom: RuntimeError reports missing or unexpected keys during load_state_dict. Cause: the class definition changed, a layer was renamed, or a wrapper such as data parallelism added prefixes. Diagnose: print sorted checkpoint keys and current model.state_dict().keys(), then compare names and shapes. Correct: instantiate the matching architecture, write a deliberate key-migration script, or load with strict=False only when the ignored keys are expected.

Symptom: loading fails with a device error on a machine without a GPU. Cause: tensors were saved from CUDA storage and the loader tried to restore them there. Diagnose: check the stack trace for CUDA availability or storage mapping. Correct: load with map_location="cpu", then move to the desired device after construction.

Symptom: predictions differ after loading although weights match. Cause: the model is still in training mode, preprocessing changed, random layers are active, or batch-normalization buffers were not saved. Diagnose: call model.eval(), compare the exact input tensor before inference, and inspect state-dictionary buffers as well as parameters. Correct: save the complete module state dictionary, freeze preprocessing versions, and test a golden input-output pair.

Symptom: an exported model works for the example input but fails for larger batches or alternate sequence lengths. Cause: tracing captured shape-specific behavior or Python-side branching. Diagnose: run export validation on the range of supported shapes. Correct: use scripting when appropriate, constrain supported shapes in the manifest, or export with a format and settings that represent dynamic axes.

Security, Performance, and Reliability

Model files deserve the same care as executable dependencies. Do not load untrusted pickle-based artifacts with torch.load. Store artifacts in controlled locations, verify digests before loading, and keep approval metadata with the model. A malicious or corrupted file can do more damage than producing a bad tensor.

Performance depends on artifact size, device movement, and load frequency. Load once during service startup rather than per request. Save checkpoints at a cadence that balances recovery value against storage and I/O cost. For large models, record whether weights are full precision, quantized, sharded, or otherwise transformed so serving code does not make false assumptions.

Reliability comes from reversibility. A model version should have an immutable identifier, metrics, schema, code revision, and rollback target. When a deployment misbehaves, operators need to know which artifact is active, which previous artifact is approved, and how to switch traffic back without retraining.

Hands-on Lab: Build a Reloadable Artifact

Prerequisites: a Python environment with PyTorch installed, a writable temporary directory, and enough familiarity with nn.Module to define a small network. No dataset is required; synthetic tensors are enough for this lab.

  1. Create a small module and run one optimizer step on synthetic data.
  2. Save a checkpoint dictionary containing model state, optimizer state, epoch, validation metric, and random number generator state.
  3. Start a fresh Python process or create fresh objects in the same script.
  4. Load the checkpoint with map_location="cpu", restore the model and optimizer, and set the model to eval for inference verification.
  5. Create a manifest with an artifact id, input feature names, validation metric, and SHA-256 digest of the weight or checkpoint file.
  6. Verify that a fixed input gives the same output before and after saving, that the digest matches, and that load_state_dict reports no missing or unexpected keys.

Verification: keep a golden input tensor and print or assert the restored output. Also assert that the manifest digest equals a freshly computed digest. Cleanup: delete temporary artifacts after the lab. In a shared registry, rollback means marking the previous artifact as active again rather than overwriting the failed one.

Assessment Exercises

  1. You add a new output layer for a different number of classes. Which keys should fail to load, and how would you prove that all earlier layers were restored correctly?
  2. A service loads a checkpoint successfully but produces unstable predictions for the same request. List three diagnostics that distinguish training-mode behavior from preprocessing drift.
  3. Design a manifest for a text classifier. What fields would you include so a future engineer can validate tokenizer compatibility and rollback safely?
  4. When would you choose TorchScript tracing, TorchScript scripting, or plain eager PyTorch loading for deployment? Explain the trade-off using model control flow and serving environment constraints.
  5. A model file in object storage has the expected name but a different digest. What should the loader do, and what operational signal should be emitted?

Summary

PyTorch model persistence is a set of deliberate boundaries. Save state dictionaries to preserve learned tensors, checkpoints to resume training, exported modules to serve in constrained runtimes, and manifests to make artifacts identifiable and verifiable. The dependable path is to reconstruct architecture explicitly, load with device-aware mapping, validate keys and shapes, test a golden prediction, record version metadata, and keep rollback artifacts immutable.