Model serialization formats (pickle, ONNX, TorchScript) in MLOps - Mini Project: Build & Apply

Model Serialization Formats: Pickle, ONNX, and TorchScript

📖 Scenario: You are working as a machine learning engineer. You have trained a simple model and now want to save it so it can be used later or shared with others. Different formats exist for saving models, such as Pickle, ONNX, and TorchScript. Each format has its own use case and benefits.In this project, you will practice saving a simple PyTorch model using these three formats step-by-step.

🎯 Goal: Build a Python script that creates a simple PyTorch model, then saves it using pickle, ONNX, and TorchScript formats. You will learn how to prepare the model, configure saving options, apply the saving commands, and finally confirm the files are created.

📋 What You'll Learn

Use PyTorch to create a simple neural network model

Save the model using Pickle format

Save the model using ONNX format

Save the model using TorchScript format

Print confirmation messages after saving each format

💡 Why This Matters

🌍 Real World

Saving machine learning models in different formats is essential for deployment, sharing, and interoperability between tools and platforms.

💼 Career

Machine learning engineers and MLOps specialists often need to serialize models efficiently and correctly for production use and collaboration.

1

Create a simple PyTorch model

Import torch and torch.nn. Create a class called SimpleModel that inherits from torch.nn.Module. Inside it, define a linear layer self.linear with input size 10 and output size 1. Implement the forward method to pass input x through self.linear. Then create an instance called model of SimpleModel.

MLOps

# Your code here

Hint

Remember to import torch and torch.nn first. Define the model class with a linear layer and a forward method. Then create the model instance.

2

Prepare a dummy input tensor for export

Create a variable called dummy_input that is a tensor of shape (1, 10) filled with random floats using torch.randn(1, 10). This will be used as input when exporting the model.

MLOps

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 1)
    def forward(self, x):
        return self.linear(x)

model = SimpleModel()

# Create dummy_input tensor with shape (1, 10)
# Your code here

Hint

Use torch.randn(1, 10) to create a random tensor with shape (1, 10).

3

Save the model using Pickle, ONNX, and TorchScript formats

Save the model using three methods:
1. Use torch.save(model, 'model_pickle.pth') to save with Pickle format.
2. Use torch.onnx.export(model, dummy_input, 'model.onnx', input_names=['input'], output_names=['output']) to save with ONNX format.
3. Use torch.jit.script(model) to create a scripted model, then save it with scripted_model.save('model_torchscript.pt').

MLOps

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 1)
    def forward(self, x):
        return self.linear(x)

model = SimpleModel()

dummy_input = torch.randn(1, 10)

# Save model using pickle
# Save model using ONNX
# Save model using TorchScript
# Your code here

Hint

Use torch.save for Pickle, torch.onnx.export for ONNX, and torch.jit.script plus .save() for TorchScript.

4

Print confirmation messages after saving each model format

Print these exact messages to confirm saving:
1. print('Pickle model saved')
2. print('ONNX model saved')
3. print('TorchScript model saved')

MLOps

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 1)
    def forward(self, x):
        return self.linear(x)

model = SimpleModel()

dummy_input = torch.randn(1, 10)

torch.save(model, 'model_pickle.pth')
torch.onnx.export(model, dummy_input, 'model.onnx', input_names=['input'], output_names=['output'])
scripted_model = torch.jit.script(model)
scripted_model.save('model_torchscript.pt')

# Print confirmation messages
# Your code here

Hint

Use three print statements with the exact messages given.

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