Convert Your PyTorch Framework to ONNX Using Wallaroo

With the rise of Machine Learning there has emerged a myriad of ML frameworks to choose from. However, when it comes to ML deployment this overabundance of frameworks can create some technical issues when it comes to putting your model into production. Due to the lack of compatibility with the different data types and operations of all these frameworks, converting or switching from one to another can be challenging. The Wallaroo platform has been designed specifically with these types of issues in mind as its agnostic design grants the ability to seamlessly click into your existing digital ecosystem and connect with everything around it so that you can take any model and then deploy it into any environment with minimal code engineering while still maintaining model testing and observability capabilities.

Wallaroo’s integrated model conversion allows you to keep your preferred framework, such as PyTorch while maintaining a standardized deployment process. The Wallaroo platform allows you to run your models in their native runtime leveraging the purpose-built model serving engine, which allows you to get the best performance out of your models. This makes your MLOps more agile, flexible, and efficient as you avoid wasting resources trying to fit your model into a deployment platform that forces you to standardize onto proprietary tools and formats. Once your models have been put into production, Wallaroo’s performance and ease of use allows for stress-free scalability, real-time predictions, and unparalleled observability.

Note: The following information has been validated only as of 11/11/2022. If you have any issues executing these procedures with similar results, please visit our documentation site for our most recent code and suggestions.

Prerequisites for Executing Wallaroo’s Model Conversion

To follow the instructions in this guide successfully, you will need a Wallaroo instance to leverage the Wallaroo platform’s self-service SDK where you can collaborate, manage, and deploy your models into production..Most tasks can be performed with either the Wallaroo SDK or Dashboard. The Wallaroo SDK can be used from the JupyterHub service included inside of Wallaroo. You can create a Wallaroo instance for yourself and try this tutorial by downloading the free Wallaroo Community Edition. Additionally, you will need the RandomForestRegressor PyTorch model pytorchbikeshare.pt, using 58 total inputs and the class BikeShareRegressor. If you are not in possession of these files they can be found via this link.

PyTorch Model Conversion

Step 1: Import the Libraries

Firstly, you need to import the libraries you will be using. In this case, we will use our PyTorch torch library and import it into the kernel. It should look like the following:

# the Pytorch libraries

# Import into this kernel

import sys

!{sys.executable} -m pip install torch

import torch

import torch.onnx

You should receive results as listed below:

Collecting torch

Downloading torch-1.12.0-cp38-cp38-manylinux1_x86_64.whl (776.3 MB)

[K     |████████████████████████████████| 776.3 MB 4.7 kB/s  eta 0:00:01     |██████                          | 146.9 MB 89.9 MB/s eta 0:00:07     |██████████████████████▍         | 544.4 MB 87.7 MB/s eta 0:00:03

[?25hRequirement already satisfied: typing-extensions in /opt/conda/lib/python3.8/site-packages (from torch) (3.7.4.3)

Installing collected packages: torch

Successfully installed torch-1.12.0

However, if our PyTorch torch library has already been imported you may see the following:

Requirement already satisfied: torch in /opt/conda/lib/python3.9/site-packages (1.13.0)

Requirement already satisfied: typing-extensions in /opt/conda/lib/python3.9/site-packages (from torch) (4.3.0)

Requirement already satisfied: nvidia-cuda-runtime-cu11==11.7.99 in /opt/conda/lib/python3.9/site-packages (from torch) (11.7.99)

Requirement already satisfied: nvidia-cudnn-cu11==8.5.0.96 in /opt/conda/lib/python3.9/site-packages (from torch) (8.5.0.96)

Requirement already satisfied: nvidia-cublas-cu11==11.10.3.66 in /opt/conda/lib/python3.9/site-packages (from torch) (11.10.3.66)

Requirement already satisfied: nvidia-cuda-nvrtc-cu11==11.7.99 in /opt/conda/lib/python3.9/site-packages (from torch) (11.7.99)

Requirement already satisfied: wheel in /opt/conda/lib/python3.9/site-packages (from nvidia-cublas-cu11==11.10.3.66->torch) (0.37.1)

Requirement already satisfied: setuptools in /opt/conda/lib/python3.9/site-packages (from nvidia-cublas-cu11==11.10.3.66->torch) (62.3.2)

Step 2: Loading the PyTorch model into a Variable

In order to load our PyTorch model into a variable, we need to first define the class of the model. So with our class as BikeShareRegressor, the defined class would be:

class BikeShareRegressor(torch.nn.Module):

def __init__(self):

super(BikeShareRegressor, self).__init__()

self.net = nn.Sequential(nn.Linear(input_size, l1),

torch.nn.ReLU(),

torch.nn.Dropout(p=dropout),

nn.BatchNorm1d(l1),

nn.Linear(l1, l2),

torch.nn.ReLU(),

torch.nn.Dropout(p=dropout),

nn.BatchNorm1d(l2),

nn.Linear(l2, output_size))

def forward(self, x):

return self.net(x)

Now, we can load our model into the variable: pytorch_tobe_converted

# load the Pytorch model

model = torch.load("./pytorch_bikesharingmodel.pt")

Step 3: Define Your Method

Next, we can define our method, here, Convert_ONNX(), using the following inputs:

PyTorchModel — the PyTorch we are converting.

modelInputs — the model input or tuple for multiple inputs.

onnxPath — The location to save the onnx file.

opset_version — The ONNX version to export to.

Input_names — Array of the model’s input names.

output_names — Array of the model’s output names.

dynamic_axes — Sets variable length axes in the format, replacing the batch_size as necessary:{‘modelInput’ : { 0 : ‘batch_size’}, ‘modelOutput’ : {0 : ‘batch_size’}}

export_params — Whether to store the trained parameter weight inside the model file. Defaults to True.

do_constant_folding — Sets whether to execute constant folding for optimization. Defaults to True.

Applying the above inputs correctly, the code should look like this:

#Function to Convert to ONNX

def Convert_ONNX():

# set the model to inference mode

model.eval()

# Export the model

torch.onnx.export(model,         # model being run

dummy_input,       # model input (or a tuple for multiple inputs)

pypath,       # where to save the model

export_params=True,  # store the trained parameter weights inside the model file

opset_version=15,    # the ONNX version to export the model to

do_constant_folding=True,  # whether to execute constant folding for optimization

input_names = ['modelInput'],   # the model's input names

output_names = ['modelOutput'], # the model's output names

dynamic_axes = {'modelInput' : {0 : 'batch_size'}, 'modelOutput' : {0 : 'batch_size'}} # variable length axes

)

print(" ")

print('Model has been converted to ONNX')

Step 4: Set Variables and Run Conversion

Here, we will use 58 as the input_size, the device value derived from torch.cuda and the ONNX exporting version will be set to 10.

The codes for the model conversion are as follows:

pypath = "pytorchbikeshare.onnx"

input_size = 58

if torch.cuda.is_available():

device = 'cuda'

else:

device = 'cpu'

onnx_opset_version = 15

# Set up some dummy input tensor for the model

dummy_input = torch.randn(1, input_size, requires_grad=True).to(device)

Convert_ONNX()

Upon successful completion, you will receive the following message:

Conclusion

Wallaroo’s agnostic deployment platform can be made to fit within any enterprise’s existing framework or ecosystem. Once in production, you can leverage all of Wallaroo’s bleeding-edge functionality to suit the needs of your enterprise’s MLOps. By taking advantage of our vetted PyTorch conversion process, you and your team will no longer be at the mercy of an all-in-one MLOps platform and can seek to exert your desired level of control over your data and models.

To further explore the value that our Model Conversions can bring to your enterprise, contact our team of experts or email us at deployML@wallaroo.ai.