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

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

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

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

Need a hint?

Use three print statements with the exact messages given.