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.

Practice

1. Which model serialization format is Python-specific and not ideal for sharing models across different platforms?

easy

A. Pickle

B. ONNX

C. TorchScript

D. JSON

Solution

1. Step 1: Understand Pickle's scope

   Pickle is a Python library that serializes Python objects but is limited to Python environments.

2. Step 2: Compare with other formats

   ONNX and TorchScript are designed for cross-platform use, unlike Pickle.

3. Final Answer:

   Pickle -> Option A

4. Quick Check:

   Python-only format = Pickle [OK]

Hint: Pickle = Python-only, others are cross-platform [OK]

Common Mistakes:

• Confusing ONNX as Python-only
• Thinking TorchScript is Python-specific
• Selecting JSON which is not a model format

2. Which of the following is the correct Python code snippet to save a PyTorch model using TorchScript?

easy

A. onnx.save(model, 'model.pt')

B. torch.save(model, 'model.pt')

C. pickle.dump(model, open('model.pt', 'wb'))

D. torch.jit.save(torch.jit.script(model), 'model.pt')

Solution

1. Step 1: Identify TorchScript saving method

   TorchScript models are saved using torch.jit.save after scripting the model with torch.jit.script.

2. Step 2: Check other options

   torch.save(model, 'model.pt') saves a PyTorch model but not as TorchScript. pickle.dump(model, open('model.pt', 'wb')) uses pickle, and onnx.save(model, 'model.pt') is invalid syntax.

3. Final Answer:

   torch.jit.save(torch.jit.script(model), 'model.pt') -> Option D

4. Quick Check:

   TorchScript save = torch.jit.save + torch.jit.script [OK]

Hint: TorchScript save needs torch.jit.script before torch.jit.save [OK]

Common Mistakes:

• Using torch.save instead of torch.jit.save
• Trying to save ONNX model with onnx.save (wrong syntax)
• Using pickle for TorchScript models

3. Given the following Python code snippet, what will be the output type of the loaded model?

import torch
import pickle

model = SomePyTorchModel()
# Save with pickle
with open('model.pkl', 'wb') as f:
    pickle.dump(model, f)

# Load model
with open('model.pkl', 'rb') as f:
    loaded_model = pickle.load(f)

print(type(loaded_model))

medium

A. <class 'torch.jit.ScriptModule'>

B. <class '__main__.SomePyTorchModel'>

C. <class 'onnx.ModelProto'>

D. TypeError

Solution

1. Step 1: Understand pickle serialization

   Pickle saves and loads the exact Python object, so the loaded model keeps the original class type.

2. Step 2: Analyze output type

   Since model was saved with pickle, loaded_model is the same class as the original model.

3. Final Answer:

   <class '__main__.SomePyTorchModel'> -> Option B

4. Quick Check:

   Pickle load returns original Python object type [OK]

Hint: Pickle load returns original Python object type [OK]

Common Mistakes:

• Confusing TorchScript or ONNX types with pickle load
• Expecting a TorchScript or ONNX model type
• Assuming a TypeError occurs on loading

4. You tried to load a model saved with TorchScript using pickle.load() and got an error. What is the most likely cause?

medium

A. TorchScript models cannot be loaded with pickle.load()

B. The model file is corrupted

C. pickle.load() requires the model to be saved as ONNX

D. TorchScript models must be loaded with torch.load()

Solution

1. Step 1: Understand serialization compatibility

   TorchScript models are saved in a special format and cannot be loaded by pickle.load(), which expects Python pickle format.

2. Step 2: Identify correct loading method

   TorchScript models should be loaded with torch.jit.load(), not pickle.load().

3. Final Answer:

   TorchScript models cannot be loaded with pickle.load() -> Option A

4. Quick Check:

   pickle.load() incompatible with TorchScript [OK]

Hint: TorchScript needs torch.jit.load(), not pickle.load() [OK]

Common Mistakes:

• Assuming torch.load() works for TorchScript
• Thinking ONNX is required for pickle.load()
• Blaming file corruption without checking method

5. You want to deploy a PyTorch model to a production environment that does not have Python installed. Which serialization format should you choose and why?

hard

A. Pickle, because it is simple and fast

B. JSON, because it stores model weights efficiently

C. TorchScript, because it can run independently of Python

D. ONNX, because it is Python-only and easy to use

Solution

1. Step 1: Identify deployment constraints

   The environment lacks Python, so the model format must run without Python dependencies.

2. Step 2: Compare serialization formats

   Pickle requires Python, ONNX is cross-platform but needs an ONNX runtime, TorchScript can run independently using PyTorch's C++ runtime.

3. Step 3: Choose best fit

   TorchScript is designed for deployment without Python, making it the best choice here.

4. Final Answer:

   TorchScript, because it can run independently of Python -> Option C

5. Quick Check:

   Deploy without Python = TorchScript [OK]

Hint: No Python? Use TorchScript for standalone deployment [OK]

Common Mistakes:

• Choosing Pickle which needs Python
• Confusing ONNX as Python-only
• Selecting JSON which is not a model format