Chapter 2: Model Formats & Serialization

WIA-AI-008 Standard • Estimated reading time: 45 minutes

2.1 Understanding Model Serialization

Serialization is the process of converting a model's in-memory representation into a format that can be stored on disk or transmitted over a network. This fundamental operation enables model persistence, versioning, and distribution. Different frameworks approach serialization with varying philosophies and technical implementations.

What Gets Serialized?

A complete model serialization typically includes:

• Architecture/Graph: The computational graph defining operations and their connections
• Parameters/Weights: Learned parameters (weights, biases) from training
• Optimizer State: (Optional) Current state of the optimizer for resuming training
• Metadata: Version info, input/output specifications, custom configurations
• Preprocessing Logic: (Optional) Data preprocessing and postprocessing code

2.2 PyTorch Serialization Formats

PyTorch provides multiple serialization mechanisms, each suited for different use cases.

torch.save() and torch.load()

The most basic serialization uses Python's pickle protocol:

import torch

# Save entire model (architecture + weights)
torch.save(model, 'model.pth')

# Save only state dict (weights only)
torch.save(model.state_dict(), 'weights.pth')

# Save checkpoint with additional info
torch.save({
    'epoch': 100,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': 0.123,
    'accuracy': 0.956
}, 'checkpoint.pth')

# Loading
model = torch.load('model.pth')  # Full model
model.load_state_dict(torch.load('weights.pth'))  # Weights only

TorchScript

TorchScript creates a serializable, optimizable representation that can run without Python:

import torch

# Tracing method (records operations)
example_input = torch.randn(1, 3, 224, 224)
traced_model = torch.jit.trace(model, example_input)
traced_model.save('model_traced.pt')

# Scripting method (analyzes code)
scripted_model = torch.jit.script(model)
scripted_model.save('model_scripted.pt')

# Load TorchScript model
loaded_model = torch.jit.load('model_traced.pt')

When to Use Each Format

Format	Use Case	Pros	Cons
torch.save (full)	Quick experiments	Simplest to use	Fragile to code changes
state_dict	Production checkpoints	Flexible, portable	Requires architecture code
TorchScript	C++ deployment, mobile	No Python dependency	Limited Python feature support

2.3 TensorFlow Serialization Formats

TensorFlow's serialization evolved significantly between TF 1.x and TF 2.x, with SavedModel becoming the standard format.

SavedModel Format

SavedModel is TensorFlow's universal serialization format:

import tensorflow as tf

# Save model
model = tf.keras.models.Sequential([...])
model.save('my_model')  # Creates directory with saved_model.pb

# Customize signatures
@tf.function(input_signature=[tf.TensorSpec(shape=[None, 224, 224, 3], dtype=tf.float32)])
def serve(x):
    return model(x)

tf.saved_model.save(model, 'my_model', signatures={'serving_default': serve})

# Load model
loaded_model = tf.saved_model.load('my_model')
# Or for Keras models:
loaded_model = tf.keras.models.load_model('my_model')

HDF5 Format (.h5)

Legacy Keras format, still widely used:

# Save to HDF5
model.save('model.h5')

# Save weights only
model.save_weights('weights.h5')

# Load
model = tf.keras.models.load_model('model.h5')

Checkpoint Format

For training checkpoints:

checkpoint = tf.train.Checkpoint(
    optimizer=optimizer,
    model=model
)

# Save checkpoint
checkpoint.save('checkpoints/ckpt')

# Restore
checkpoint.restore('checkpoints/ckpt-10')

2.4 JAX Model Serialization

JAX doesn't have built-in serialization since it's a functional framework. The community uses several approaches:

Using Pickle

import pickle
import jax.numpy as jnp

# Save parameters
params = {'w': jnp.array([1., 2., 3.]), 'b': jnp.array(0.5)}
with open('params.pkl', 'wb') as f:
    pickle.dump(params, f)

# Load parameters
with open('params.pkl', 'rb') as f:
    loaded_params = pickle.load(f)

Using MessagePack/Orbax

from flax import serialization

# Serialize to bytes
bytes_output = serialization.to_bytes(params)

# Save to file
with open('params.msgpack', 'wb') as f:
    f.write(bytes_output)

# Load
with open('params.msgpack', 'rb') as f:
    loaded_params = serialization.from_bytes(params, f.read())

2.5 ONNX Format Deep Dive

ONNX (Open Neural Network Exchange) is a framework-agnostic format that serves as a universal intermediate representation.

ONNX File Structure

An ONNX file (.onnx) contains:

• Graph: Computational graph with nodes (operations) and edges (tensors)
• Operators: Standardized operations (Conv, MatMul, ReLU, etc.)
• Tensors: Weights stored as Protocol Buffer messages
• Metadata: Model version, producer, opset version

import onnx

# Load and inspect ONNX model
model = onnx.load('model.onnx')

print(f"IR Version: {model.ir_version}")
print(f"Producer: {model.producer_name}")
print(f"Opset Version: {model.opset_import[0].version}")

# Check model validity
onnx.checker.check_model(model)

# Print graph info
graph = model.graph
print(f"Inputs: {[i.name for i in graph.input]}")
print(f"Outputs: {[o.name for o in graph.output]}")
print(f"Nodes: {len(graph.node)}")

Export to ONNX

From PyTorch:

import torch.onnx

dummy_input = torch.randn(1, 3, 224, 224)

torch.onnx.export(
    model,                          # Model
    dummy_input,                    # Example input
    "resnet50.onnx",               # Output file
    export_params=True,             # Store weights
    opset_version=14,              # ONNX version
    do_constant_folding=True,      # Optimize constant folding
    input_names=['input'],         # Input names
    output_names=['output'],       # Output names
    dynamic_axes={                  # Variable dimensions
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)

From TensorFlow:

import tf2onnx

# Convert SavedModel to ONNX
python -m tf2onnx.convert \
    --saved-model my_model \
    --output model.onnx \
    --opset 14

# Or programmatically
import tensorflow as tf
import tf2onnx

spec = (tf.TensorSpec((None, 224, 224, 3), tf.float32, name="input"),)
output_path = "model.onnx"

model_proto, _ = tf2onnx.convert.from_keras(model, input_signature=spec, opset=14)
with open(output_path, "wb") as f:
    f.write(model_proto.SerializeToString())

2.6 TensorFlow Lite Format

TensorFlow Lite (.tflite) is optimized for mobile and edge devices:

import tensorflow as tf

# Convert to TFLite
converter = tf.lite.TFLiteConverter.from_saved_model('my_model')

# Optimization options
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_types = [tf.float16]  # FP16 quantization

tflite_model = converter.convert()

# Save
with open('model.tflite', 'wb') as f:
    f.write(tflite_model)

# Inference
interpreter = tf.lite.Interpreter(model_path='model.tflite')
interpreter.allocate_tensors()

input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

interpreter.set_tensor(input_details[0]['index'], input_data)
interpreter.invoke()
output_data = interpreter.get_tensor(output_details[0]['index'])

2.7 CoreML Format

CoreML is Apple's format for iOS/macOS deployment:

import coremltools as ct

# Convert PyTorch to CoreML
example_input = torch.rand(1, 3, 224, 224)
traced_model = torch.jit.trace(model, example_input)

coreml_model = ct.convert(
    traced_model,
    inputs=[ct.TensorType(shape=(1, 3, 224, 224))]
)

coreml_model.save("MyModel.mlmodel")

# Add metadata
coreml_model.author = 'WIA Team'
coreml_model.license = 'Apache 2.0'
coreml_model.short_description = 'Image classification model'

2.8 Model Format Comparison

Format	Framework	File Ext	Primary Use	Size
PyTorch State Dict	PyTorch	.pth, .pt	Training/checkpoints	Medium
TorchScript	PyTorch	.pt	Production/C++	Medium
SavedModel	TensorFlow	Directory	Production/serving	Large
HDF5	Keras	.h5	Legacy Keras	Small
ONNX	Universal	.onnx	Interoperability	Medium
TFLite	TensorFlow	.tflite	Mobile/edge	Very Small
CoreML	Apple	.mlmodel	iOS/macOS	Small

2.9 Versioning and Compatibility

Model formats evolve over time. Ensuring compatibility requires careful version management:

Semantic Versioning for Models

# Model version format: MAJOR.MINOR.PATCH
# Example: 2.1.3

# MAJOR: Incompatible API changes (input/output shape changes)
# MINOR: Backward-compatible functionality (new features, improved accuracy)
# PATCH: Backward-compatible bug fixes (numerical fixes, no architecture change)

metadata = {
    'model_name': 'sentiment-classifier',
    'version': '2.1.3',
    'framework': 'pytorch',
    'framework_version': '2.0.1',
    'onnx_opset': 14,
    'created_at': '2025-01-15T10:30:00Z'
}

2.10 Best Practices for Serialization

1. Always Save Metadata

# Bad: Only weights
torch.save(model.state_dict(), 'model.pth')

# Good: Include context
torch.save({
    'model_state_dict': model.state_dict(),
    'model_name': 'resnet50',
    'version': '1.0.0',
    'input_shape': [1, 3, 224, 224],
    'num_classes': 1000,
    'accuracy': 0.923,
    'framework_version': torch.__version__,
    'created_at': datetime.now().isoformat()
}, 'model_checkpoint.pth')

2. Use Framework-Agnostic Formats for Distribution

Prefer ONNX for sharing models that may be used in different frameworks or platforms.

3. Validate After Serialization

import numpy as np

# Generate test input
test_input = torch.randn(5, 3, 224, 224)

# Original model output
original_output = model(test_input)

# Save and load
torch.save(model.state_dict(), 'model.pth')
loaded_model = ResNet50()
loaded_model.load_state_dict(torch.load('model.pth'))
loaded_model.eval()

# Loaded model output
loaded_output = loaded_model(test_input)

# Validate numerical equivalence
assert torch.allclose(original_output, loaded_output, rtol=1e-5)
print("✓ Serialization validation passed")

4. Include Preprocessing Information

preprocessing_info = {
    'mean': [0.485, 0.456, 0.406],
    'std': [0.229, 0.224, 0.225],
    'resize': 256,
    'crop': 224,
    'normalization': 'imagenet'
}