Saving and Loading

# Save (model weights only)
torch.save(model.state_dict(), 'model.pt')

# Load
model = MyModel()
model.load_state_dict(torch.load('model.pt', weights_only=True))
model.eval()

# Save
checkpoint = {
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'scheduler_state_dict': scheduler.state_dict(),
    'loss': loss,
    'config': {'lr': 1e-3, 'batch_size': 32},
}
torch.save(checkpoint, 'checkpoint.pt')

# Load and resume
checkpoint = torch.load('checkpoint.pt', map_location=device, weights_only=True)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
start_epoch = checkpoint['epoch'] + 1

# strict=True (default): every key must match exactly
model.load_state_dict(state, strict=True)

# strict=False: allow missing or extra keys (e.g., fine-tuning, partial load)
missing, unexpected = model.load_state_dict(state, strict=False)
print('Missing keys:', missing)
print('Unexpected keys:', unexpected)

# assign=False (default): copies data into existing params in-place
# assign=True (PyTorch ≥ 2.1): assigns new tensors — needed for meta device
model.load_state_dict(state, assign=True)

# Load GPU checkpoint on CPU
state = torch.load('model.pt', map_location='cpu')

# Load on specific GPU
state = torch.load('model.pt', map_location='cuda:1')

# Dynamic device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
state = torch.load('model.pt', map_location=device)

# Remap specific devices
state = torch.load('model.pt', map_location={'cuda:0': 'cuda:1'})

# Save any Python object that pickle can handle
torch.save(obj, 'file.pt')

# Load it back
obj = torch.load('file.pt', weights_only=True)   # tensors/dicts/lists only
obj = torch.load('file.pt', weights_only=False)  # full pickle (trusted source)

# Save multiple tensors in a dict
torch.save({'x': x, 'y': y, 'z': z}, 'tensors.pt')
data = torch.load('tensors.pt', weights_only=True)

# Save to a file-like object (BytesIO, network stream, etc.)
import io
buf = io.BytesIO()
torch.save(model.state_dict(), buf)
buf.seek(0)
state = torch.load(buf, weights_only=True)

# Inspect the state_dict
sd = model.state_dict()
for name, tensor in sd.items():
    print(name, tensor.shape, tensor.dtype)

# Optimizer state_dict
opt_sd = optimizer.state_dict()
# Contains: 'state' (per-param moments) and 'param_groups' (hyperparams)

# Manually modify a value before loading (e.g., resize an embedding)
sd['embedding.weight'] = new_embedding_weights
model.load_state_dict(sd)

pretrained = torch.load('pretrained.pt', weights_only=True)
model_sd = model.state_dict()

# Filter out keys with mismatched shapes
filtered = {
    k: v for k, v in pretrained.items()
    if k in model_sd and v.shape == model_sd[k].shape
}

model_sd.update(filtered)
model.load_state_dict(model_sd)

# If your state_dict contains custom classes alongside tensors
import torch.serialization

torch.serialization.add_safe_globals([MyCustomClass])
state = torch.load('checkpoint.pt', weights_only=True)

# Method 1: tracing — record operations on example inputs
example_input = torch.rand(1, 3, 224, 224)
traced = torch.jit.trace(model.eval(), example_input)
traced.save('model_traced.pt')

# Method 2: scripting — parse Python source (supports control flow)
scripted = torch.jit.script(model)
scripted.save('model_scripted.pt')

# Load (no class definition needed)
loaded = torch.jit.load('model_traced.pt')
output = loaded(input_tensor)

# Recommended (PyTorch >= 2.5): dynamo=True — torch.export-based exporter
onnx_program = torch.onnx.export(
    model,
    (example_input,),
    dynamo=True,
    dynamic_shapes={'x': {0: 'batch_size'}},   # optional, torch.export-style
)
onnx_program.save('model.onnx')

# Legacy TorchScript-based exporter (dynamo=False, the old default)
import torch.onnx

torch.onnx.export(
    model,
    args=example_input,           # or tuple of inputs
    f='model.onnx',
    export_params=True,           # include weights
    opset_version=17,             # ONNX opset (use recent stable)
    do_constant_folding=True,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input':  {0: 'batch_size'},
        'output': {0: 'batch_size'},
    },
)

# Verify the exported model
import onnx
onnx_model = onnx.load('model.onnx')
onnx.checker.check_model(onnx_model)

import torch.export

# Export with strict symbolic tracing
exported = torch.export.export(model, args=(example_input,))
exported.module()(example_input)   # run exported graph

# Save / load ExportedProgram
torch.export.save(exported, 'exported.pt2')
loaded = torch.export.load('exported.pt2')

# With DDP, save only on rank 0
if dist.get_rank() == 0:
    torch.save(model.module.state_dict(), 'model.pt')   # model.module unwraps DDP

# Load on each rank
model.load_state_dict(torch.load('model.pt', map_location={'cuda:0': f'cuda:{rank}'}))

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import StateDictType, FullStateDictConfig

# Gather full state_dict on rank 0
with FSDP.state_dict_type(
    model,
    StateDictType.FULL_STATE_DICT,
    FullStateDictConfig(offload_to_cpu=True, rank0_only=True),
):
    state = model.state_dict()
    if dist.get_rank() == 0:
        torch.save(state, 'model_fsdp.pt')

# 1. Always include model architecture version in checkpoint
checkpoint = {
    'model_state_dict': model.state_dict(),
    'arch': 'resnet50_v2',
    'pytorch_version': torch.__version__,
}

# 2. Use atomic writes to avoid corruption
import os, tempfile

def safe_save(obj, path):
    dir_ = os.path.dirname(os.path.abspath(path))
    with tempfile.NamedTemporaryFile(dir=dir_, delete=False) as tmp:
        torch.save(obj, tmp.name)
    os.replace(tmp.name, path)   # atomic on POSIX

# 3. Keep the last K checkpoints
import glob
ckpts = sorted(glob.glob('ckpt_epoch_*.pt'), key=os.path.getmtime)
for old in ckpts[:-3]:   # keep last 3
    os.remove(old)

for state in optimizer.state.values():
    for k, v in state.items():
        if isinstance(v, torch.Tensor):
            state[k] = v.to(device)