Model Serialization

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'
# 'bias'
# }

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
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

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
# Quantized

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
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