Model serialization formats (pickle, ONNX, TorchScript) in MLOps - Commands & Configuration

Introduction

Saving machine learning models lets you reuse them later without retraining. Different formats like Pickle, ONNX, and TorchScript help store models so you can load and use them easily in various environments.

When you want to save a simple Python model quickly for later use in the same environment.

When you need to share a model between different programming languages or frameworks.

When you want to optimize a PyTorch model for faster inference and deployment.

When you want to move a model from training on your laptop to running on a server without retraining.

When you want to convert a model to a standard format for compatibility with many tools.

Commands

This command trains a simple logistic regression model and saves it using Pickle format to 'model.pkl'. Pickle stores Python objects in a file.

Terminal

python -c "import pickle; from sklearn.linear_model import LogisticRegression; model = LogisticRegression(); model.fit([[0,0],[1,1]], [0,1]); with open('model.pkl', 'wb') as f: pickle.dump(model, f)"

Expected Output

No output (command runs silently)

This loads the Pickle model from 'model.pkl' and runs a prediction on new data to verify it works.

Terminal

python -c "import pickle; with open('model.pkl', 'rb') as f: model = pickle.load(f); print(model.predict([[2,2]]))"

Expected Output

[1]

This creates a simple PyTorch model, converts it to TorchScript format for optimized saving, and saves it as 'model.pt'.

Terminal

python -c "import torch; import torch.nn as nn; class SimpleModel(nn.Module): def __init__(self): super().__init__(); self.linear = nn.Linear(2,1); def forward(self, x): return self.linear(x); model = SimpleModel(); scripted_model = torch.jit.script(model); scripted_model.save('model.pt')"

Expected Output

No output (command runs silently)

This loads the TorchScript model from 'model.pt' and runs a prediction to check it loads correctly.

Terminal

python -c "import torch; scripted_model = torch.jit.load('model.pt'); print(scripted_model(torch.tensor([[1.0, 2.0]])))"

Expected Output

tensor([[0.0]], grad_fn=<AddmmBackward0>)

This exports a PyTorch model to ONNX format as 'model.onnx' which is a standard format for interoperability.

Terminal

python -c "import torch; import torch.onnx; model = torch.nn.Linear(2, 1); dummy_input = torch.randn(1, 2); torch.onnx.export(model, dummy_input, 'model.onnx', input_names=['input'], output_names=['output'])"

Expected Output

No output (command runs silently)

Key Concept

If you remember nothing else from this pattern, remember: choose the model format based on where and how you want to use the model later.

Code Example

MLOps

import pickle
from sklearn.linear_model import LogisticRegression
import torch
import torch.nn as nn
import torch.onnx

# Train and save Pickle model
model = LogisticRegression()
model.fit([[0,0],[1,1]], [0,1])
with open('model.pkl', 'wb') as f:
    pickle.dump(model, f)

# Load Pickle model and predict
with open('model.pkl', 'rb') as f:
    loaded_model = pickle.load(f)
print('Pickle prediction:', loaded_model.predict([[2,2]]))

# Define and save TorchScript model
class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(2,1)
    def forward(self, x):
        return self.linear(x)

model_ts = SimpleModel()
scripted_model = torch.jit.script(model_ts)
scripted_model.save('model.pt')

# Load TorchScript model and predict
loaded_ts = torch.jit.load('model.pt')
print('TorchScript prediction:', loaded_ts(torch.tensor([[1.0, 2.0]])))

# Export to ONNX
model_onnx = nn.Linear(2,1)
dummy_input = torch.randn(1,2)
torch.onnx.export(model_onnx, dummy_input, 'model.onnx', input_names=['input'], output_names=['output'])
print('ONNX model exported')

Common Mistakes

Using Pickle to save models for use in different programming languages.

Pickle is Python-specific and won't work outside Python environments.

Use ONNX for cross-language compatibility.

Trying to load a TorchScript model with regular PyTorch load functions.

TorchScript models require torch.jit.load to load properly.

Always use torch.jit.load for TorchScript models.

Not providing dummy input when exporting to ONNX.

ONNX export needs example input to trace the model graph.

Always provide a dummy input tensor matching model input shape.

Summary

Use Pickle to save and load Python models quickly in the same environment.

Use TorchScript to optimize and save PyTorch models for efficient deployment.

Use ONNX to export models for use across different frameworks and languages.