Model Serialization: Saving and Loading Models

Overview

Saving and loading models is critical for: - Checkpointing training progress - Deployment to production - Sharing models (HuggingFace Hub) - Reproducibility across systems

Different serialization formats have trade-offs in safety, compatibility, and size.

---

PyTorch Native: torch.save / torch.load

Saving State Dictionary

import torch
import torch.nn as nn

# Create and train model
model = nn.Linear(10, 5)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# ... training ...

# Save state dictionary
torch.save(model.state_dict(), 'model.pth')

# state_dict contains:
# {
#     'weight': tensor([...]),
#     'bias': tensor([...])
# }

Loading State Dictionary

# Create fresh model
model = nn.Linear(10, 5)

# Load weights
model.load_state_dict(torch.load('model.pth'))
model.eval()  # Set to evaluation mode

# Use for inference
with torch.no_grad():
    output = model(input_data)

Saving Complete Checkpoint

# Save everything needed to resume training
checkpoint = {
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': loss,
    'config': model_config,
}

torch.save(checkpoint, f'checkpoint_epoch_{epoch}.pth')

# Later: resume training
checkpoint = torch.load('checkpoint_epoch_10.pth')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']

# Continue training from same point

Entire Model (Discouraged)

# Save entire model
torch.save(model, 'model.pth')  # Pickles whole model

# Load
model = torch.load('model.pth')

# Problems:
# 1. Large file (includes code)
# 2. Compatibility issues (code changes)
# 3. Security risk (arbitrary code execution)
# 4. Python version dependencies

# ✓ Always use state_dict instead

---

State Dictionary Structure

Named Parameters

import torch.nn as nn

model = nn.Sequential(
    nn.Linear(10, 20),
    nn.ReLU(),
    nn.Linear(20, 5)
)

# Get state dict
state = model.state_dict()

print(state.keys())
# odict_keys([
#     '0.weight', '0.bias',      # First Linear layer
#     '2.weight', '2.bias'       # Third Linear layer (ReLU has no params)
# ])

for name, param in state.items():
    print(f"{name}: {param.shape}")
    # 0.weight: torch.Size([20, 10])
    # 0.bias: torch.Size([20])
    # 2.weight: torch.Size([5, 20])
    # 2.bias: torch.Size([5])

Custom Module State

class CustomModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(10, 5)
        self.scale = nn.Parameter(torch.ones(5))  # Trainable
        self.constant = torch.tensor(0.5)  # Not in state_dict

    def forward(self, x):
        return self.fc(x) * self.scale

model = CustomModel()
state = model.state_dict()

print(state.keys())
# odict_keys(['fc.weight', 'fc.bias', 'scale'])
# Note: 'constant' not included (not Parameter)

State Dict on Different Device

# Save on GPU
model = model.cuda()
torch.save(model.state_dict(), 'model.pth')

# Load on CPU
device = 'cpu'
state_dict = torch.load('model.pth', map_location=device)
model.load_state_dict(state_dict)

# Or with explicit mapping
state_dict = torch.load('model.pth', map_location={
    'cuda:0': 'cpu',      # GPU 0 → CPU
    'cuda:1': 'cuda:2',   # GPU 1 → GPU 2
})

---

SafeTensors Format

Why SafeTensors?

# Problems with torch.save:
# 1. Uses pickle (security risk)
# 2. Slow loading
# 3. Large file size
# 4. Requires full model code to load

# SafeTensors:
# 1. Safe format (just data, no code)
# 2. Fast loading (mmap support)
# 3. Smaller files (no metadata overhead)
# 4. Language-agnostic (used by JAX, TF too)

Saving and Loading

from safetensors.torch import save_file, load_file

# Save
state_dict = model.state_dict()
save_file(state_dict, 'model.safetensors')

# Load
state_dict = load_file('model.safetensors')
model.load_state_dict(state_dict)

# Advantages:
# - File is ~2x smaller
# - Loading 10x faster
# - Safe (no arbitrary code execution)
# - Compatible with HuggingFace Hub

Lazy Loading (Memory Efficient)

from safetensors.torch import load_file

# Load only specific keys (without loading full model)
state_dict = load_file('model.safetensors', 
                       keys=['fc.weight', 'fc.bias'])

# Or iterate without loading all to memory
with safe_open('model.safetensors', framework='pt') as f:
    for key in f.keys():
        tensor = f.get_tensor(key)
        print(f"Loaded {key}")

---

HuggingFace Model Hub

Publishing Models

from transformers import AutoModel, AutoTokenizer

# Load model
model = AutoModel.from_pretrained('bert-base-uncased')

# Push to Hub
model.push_to_hub("username/my-bert-model")

# Push with config
model.push_to_hub(
    repo_id="username/my-bert-model",
    commit_message="Add my custom BERT variant",
    private=False
)

Downloading Models

from transformers import AutoModel

# Download and load
model = AutoModel.from_pretrained('username/my-bert-model')

# Download to cache directory
model = AutoModel.from_pretrained(
    'username/my-bert-model',
    cache_dir='/path/to/cache'
)

# Load specific revision
model = AutoModel.from_pretrained(
    'username/my-bert-model',
    revision='main'  # or specific branch/tag
)

Custom Config and Tokenizer

from transformers import BertConfig, BertModel, BertTokenizer

# Create model with custom config
config = BertConfig(
    vocab_size=50265,
    hidden_size=512,
    num_hidden_layers=6,
    num_attention_heads=8,
)

model = BertModel(config)

# Save with config
model.save_pretrained('my-model')
config.save_pretrained('my-model')

# Load entire directory
model = BertModel.from_pretrained('my-model')
config = BertConfig.from_pretrained('my-model')
tokenizer = BertTokenizer.from_pretrained('my-model')

---

Checkpoint Management

Saving Best Model

import torch
import torch.nn as nn

class CheckpointManager:
    def __init__(self, model, optimizer, save_dir='./checkpoints'):
        self.model = model
        self.optimizer = optimizer
        self.save_dir = Path(save_dir)
        self.save_dir.mkdir(exist_ok=True)
        self.best_loss = float('inf')

    def save_checkpoint(self, epoch, loss, is_best=False):
        """Save checkpoint."""
        checkpoint = {
            'epoch': epoch,
            'model_state_dict': self.model.state_dict(),
            'optimizer_state_dict': self.optimizer.state_dict(),
            'loss': loss,
        }

        # Save latest
        latest_path = self.save_dir / 'latest.pth'
        torch.save(checkpoint, latest_path)

        # Save if best
        if is_best:
            best_path = self.save_dir / 'best.pth'
            torch.save(checkpoint, best_path)
            self.best_loss = loss

        # Save periodic
        if epoch % 10 == 0:
            period_path = self.save_dir / f'checkpoint_epoch_{epoch}.pth'
            torch.save(checkpoint, period_path)

    def load_checkpoint(self, path):
        """Load checkpoint."""
        checkpoint = torch.load(path)
        self.model.load_state_dict(checkpoint['model_state_dict'])
        self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        return checkpoint['epoch'], checkpoint['loss']

# Usage
manager = CheckpointManager(model, optimizer)

for epoch in range(100):
    train_loss = train_epoch(model, optimizer, train_loader)
    val_loss = validate(model, val_loader)

    is_best = val_loss < manager.best_loss
    manager.save_checkpoint(epoch, val_loss, is_best=is_best)

Loading Last Checkpoint

import os
from pathlib import Path

def find_latest_checkpoint(checkpoint_dir):
    """Find most recent checkpoint."""
    checkpoint_dir = Path(checkpoint_dir)

    # First try 'latest.pth'
    latest = checkpoint_dir / 'latest.pth'
    if latest.exists():
        return latest

    # Otherwise find newest by modification time
    checkpoints = list(checkpoint_dir.glob('checkpoint_*.pth'))
    if not checkpoints:
        return None

    return max(checkpoints, key=lambda p: p.stat().st_mtime)

# Usage
latest_ckpt = find_latest_checkpoint('./checkpoints')
if latest_ckpt:
    checkpoint = torch.load(latest_ckpt)
    model.load_state_dict(checkpoint['model_state_dict'])
    print(f"Resumed from {latest_ckpt}")
else:
    print("No checkpoint found, training from scratch")

---

Cross-Framework Serialization

ONNX Format (Inference)

import torch
import torch.onnx

# Create model
model = nn.Linear(10, 5)
model.eval()

# Export to ONNX
dummy_input = torch.randn(1, 10)
torch.onnx.export(
    model,
    dummy_input,
    'model.onnx',
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={'input': {0: 'batch_size'}}
)

# Load and use with ONNX Runtime
import onnxruntime

sess = onnxruntime.InferenceSession('model.onnx')
output = sess.run(None, {'input': input_data.numpy()})

TorchScript for Deployment

# Create model
model = nn.Linear(10, 5)

# Script (trace-based)
dummy_input = torch.randn(1, 10)
traced_model = torch.jit.trace(model, dummy_input)
traced_model.save('model.pt')

# Script (compile-based)
scripted_model = torch.jit.script(model)
scripted_model.save('model.pt')

# Load and use
loaded_model = torch.jit.load('model.pt')
output = loaded_model(input_data)

---

Storage and Bandwidth Optimization

Quantization Before Saving

import torch
import torch.quantization as quantization

# Create model
model = nn.Sequential(
    nn.Linear(10, 20),
    nn.ReLU(),
    nn.Linear(20, 5)
)

# Quantize to int8 (4x smaller)
model.qconfig = quantization.get_default_qconfig('fbgemm')
quantization.prepare(model, inplace=True)
quantization.convert(model, inplace=True)

# Save quantized model
torch.save(model.state_dict(), 'model_quantized.pth')

# File size comparison:
# Original: ~4MB
# Quantized: ~1MB

Compression Techniques

# 1. Weight pruning (remove small weights)
def prune_weights(state_dict, threshold=1e-6):
    pruned = {}
    for key, tensor in state_dict.items():
        mask = torch.abs(tensor) > threshold
        pruned[key] = tensor * mask
    return pruned

pruned_state = prune_weights(model.state_dict())
torch.save(pruned_state, 'model_pruned.pth')

# 2. Shared weights (same tensor multiple places)
# 3. Compression libraries (gzip, brotli)
import gzip

state_dict = torch.load('model.pth')
with gzip.open('model.pth.gz', 'wb') as f:
    torch.save(state_dict, f)

# 4. Different precision (float32 → float16)
state_half = {}
for key, tensor in model.state_dict().items():
    state_half[key] = tensor.half()

torch.save(state_half, 'model_fp16.pth')

---

Practical: Complete Training with Checkpoints

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

class Trainer:
    def __init__(self, model, optimizer, criterion, checkpoint_dir='./checkpoints'):
        self.model = model
        self.optimizer = optimizer
        self.criterion = criterion
        self.checkpoint_dir = Path(checkpoint_dir)
        self.checkpoint_dir.mkdir(exist_ok=True)

    def train_epoch(self, train_loader):
        self.model.train()
        total_loss = 0

        for x, y in train_loader:
            pred = self.model(x)
            loss = self.criterion(pred, y)

            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()

            total_loss += loss.item()

        return total_loss / len(train_loader)

    def evaluate(self, val_loader):
        self.model.eval()
        total_loss = 0

        with torch.no_grad():
            for x, y in val_loader:
                pred = self.model(x)
                loss = self.criterion(pred, y)
                total_loss += loss.item()

        return total_loss / len(val_loader)

    def save_checkpoint(self, epoch, val_loss, is_best=False):
        """Save checkpoint."""
        checkpoint = {
            'epoch': epoch,
            'model_state_dict': self.model.state_dict(),
            'optimizer_state_dict': self.optimizer.state_dict(),
            'val_loss': val_loss,
        }

        # Save latest
        torch.save(checkpoint, self.checkpoint_dir / 'latest.pth')

        # Save best
        if is_best:
            torch.save(checkpoint, self.checkpoint_dir / 'best.pth')

    def load_checkpoint(self, checkpoint_name='latest.pth'):
        """Load checkpoint."""
        path = self.checkpoint_dir / checkpoint_name
        if not path.exists():
            return 0

        checkpoint = torch.load(path)
        self.model.load_state_dict(checkpoint['model_state_dict'])
        self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        return checkpoint['epoch']

    def train(self, train_loader, val_loader, num_epochs=100):
        """Complete training loop with checkpointing."""
        start_epoch = self.load_checkpoint()
        best_val_loss = float('inf')

        for epoch in range(start_epoch, num_epochs):
            train_loss = self.train_epoch(train_loader)
            val_loss = self.evaluate(val_loader)

            is_best = val_loss < best_val_loss
            if is_best:
                best_val_loss = val_loss

            self.save_checkpoint(epoch, val_loss, is_best=is_best)

            print(f"Epoch {epoch}: train_loss={train_loss:.4f}, "
                  f"val_loss={val_loss:.4f}, best={best_val_loss:.4f}")

# Usage
model = nn.Linear(10, 5)
optimizer = optim.SGD(model.parameters(), lr=0.01)
criterion = nn.MSELoss()

trainer = Trainer(model, optimizer, criterion)
trainer.train(train_loader, val_loader, num_epochs=100)

# Later: resume training
trainer.load_checkpoint()
trainer.train(train_loader, val_loader, num_epochs=200)

---

Summary: Serialization Formats

Format	Use Case	Size	Speed	Safety
pickle	Small models	Large	Slow	❌ Unsafe
state_dict	Training checkpoints	Medium	Medium	✅ Safe
safetensors	Sharing/deployment	Small	Fast	✅ Safe
ONNX	Cross-framework inference	Medium	Fast	✅ Safe
TorchScript	Deployment without code	Medium	Fast	✅ Safe
HuggingFace	Community models	Variable	Medium	✅ Safe

---

Related Topics

• 02 Autograd Implementation - Understanding state_dict contents
• 05 Custom Operators - Saving custom operations
• 06 Inference Optimization Patterns - Deployment considerations