Save / Load

import torch
import torch.nn as nn
from torch.optim import AdamW
from pathlib import Path

# 1) Model + optimizer
class MLP(nn.Module):
    def __init__(self, d_in=784, d_hidden=256, d_out=10):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(d_in, d_hidden), nn.ReLU(),
            nn.Linear(d_hidden, d_hidden), nn.ReLU(),
            nn.Linear(d_hidden, d_out),
        )
    def forward(self, x): return self.net(x)

model = MLP()
opt   = AdamW(model.parameters(), lr=3e-4)

# 2) Save state_dict — the recommended way
torch.save(model.state_dict(), 'model.pt')
# state_dict is just a Python OrderedDict[str, Tensor].
# Smaller, more portable, less coupled to source code than pickling the module.

# 3) Load state_dict
model2 = MLP()                              # build the SAME architecture first
model2.load_state_dict(torch.load('model.pt', map_location='cpu'))
model2.eval()                                # switch off dropout/batchnorm-update

# 4) Strict vs lax loading
model2.load_state_dict(state, strict=False) # tolerates missing or extra keys (returns the diff)

# When you've added a new layer and want to load partial weights:
incompat = model2.load_state_dict(old_state, strict=False)
print('missing:', incompat.missing_keys)
print('unexpected:', incompat.unexpected_keys)

# 5) Full checkpoint — model + optimizer + epoch + RNG
ckpt = {
    'model':    model.state_dict(),
    'optimizer': opt.state_dict(),
    'epoch':    epoch,
    'best_val_loss': best_val_loss,
    'rng_state':    torch.get_rng_state(),
    'cuda_rng_state': torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None,
}
torch.save(ckpt, 'ckpt-epoch-12.pt')

# 6) Resume training
ckpt = torch.load('ckpt-epoch-12.pt', map_location=device)
model.load_state_dict(ckpt['model'])
opt.load_state_dict(ckpt['optimizer'])
start_epoch     = ckpt['epoch'] + 1
best_val_loss   = ckpt['best_val_loss']
torch.set_rng_state(ckpt['rng_state'])
if ckpt['cuda_rng_state'] is not None:
    torch.cuda.set_rng_state_all(ckpt['cuda_rng_state'])

# 7) Save-best pattern
best_val = float('inf')
for epoch in range(epochs):
    train_one_epoch(model, opt, train_loader)
    val_loss = evaluate(model, val_loader)
    if val_loss < best_val:
        best_val = val_loss
        torch.save(model.state_dict(), 'best.pt')
    if (epoch + 1) % 5 == 0:
        torch.save(model.state_dict(), f'epoch-{epoch+1:03d}.pt')

# 8) Save with safetensors — recommended for distribution
# pip install safetensors
from safetensors.torch import save_file, load_file
save_file(model.state_dict(), 'model.safetensors')
state = load_file('model.safetensors', device='cpu')
model.load_state_dict(state)
# Why: safer (no pickle code execution), faster, memory-mapped — the standard format for HF models.

# 9) torch.save vs pickling the whole model
# torch.save(model, 'whole.pt') — saves the MODULE, not just the weights.
# Pros: don't need the class to load (sort of)
# Cons: tied to internal class paths, breaks on refactors, weight_only=True
#       can't load it. AVOID for long-term storage.
#
# Use it only when you'll load on the same code base, same version.

# 10) Distributed Data Parallel — strip the 'module.' prefix
# Saving with DDP: save model.module.state_dict() so non-DDP code can load it.
if hasattr(model, 'module'):
    torch.save(model.module.state_dict(), 'best.pt')
else:
    torch.save(model.state_dict(), 'best.pt')

# Loading a DDP-saved checkpoint into a plain model:
state = torch.load('best.pt', map_location='cpu')
state = { k.replace('module.', ''): v for k, v in state.items() }
model.load_state_dict(state)

# 11) TorchScript — bundle architecture + weights for deployment
# Trace (works for static control flow)
example = torch.randn(1, 784)
traced = torch.jit.trace(model.eval(), example)
traced.save('model.ts.pt')
# Load anywhere without the Python class:
loaded = torch.jit.load('model.ts.pt')
out = loaded(torch.randn(1, 784))

# Script (works for dynamic control flow)
scripted = torch.jit.script(model.eval())
scripted.save('model_scripted.ts.pt')

# 12) ONNX — interop with other runtimes (TensorRT, ONNX Runtime, browsers)
torch.onnx.export(
    model.eval(),
    example,
    'model.onnx',
    input_names=['x'],
    output_names=['logits'],
    dynamic_axes={'x': {0: 'batch'}, 'logits': {0: 'batch'}},
    opset_version=17,
)

# Validate
import onnx, onnxruntime as ort
onnx.checker.check_model(onnx.load('model.onnx'))
sess = ort.InferenceSession('model.onnx')
ort_out = sess.run(['logits'], {'x': example.numpy()})

# 13) Quantize for smaller / faster inference
from torch.ao.quantization import quantize_dynamic
q = quantize_dynamic(model.eval(), {nn.Linear}, dtype=torch.qint8)
torch.save(q.state_dict(), 'model.qint8.pt')

# 14) Versioning + metadata
import json, time
meta = {
    'name':       'mlp',
    'version':    '1.2.3',
    'created':    int(time.time()),
    'pytorch':    torch.__version__,
    'input_shape': [1, 784],
    'mean':       [0.1307],
    'std':        [0.3081],
}
Path('model.json').write_text(json.dumps(meta, indent=2))
# Ship model.safetensors + model.json together; reproducibility is half the battle.

# 15) Common bugs
#   • Saving with one device, loading on another → set map_location='cpu' or 'cuda'
#   • Loading state_dict into a different architecture → strict=False + manual reconcile
#   • Forgetting model.eval() before inference → dropout still active, batchnorm uses batch stats
#   • Trace exported with dynamic shapes → falls back to fixed; use scripting or set dynamic_axes
#   • Saving the whole model (pickle) then refactoring class → can't load
#   • Optimizer state lost on resume → loss spikes; always save and restore both
#   • DDP 'module.' prefix → KeyError on non-DDP load

torch.save(model.state_dict(), 'model.pt')
# Restore
model = MLP()
model.load_state_dict(torch.load('model.pt', map_location='cpu'))
model.eval()