What every ML/AI developer should know about ONNX

The Open Neural Network Exchange Format (ONNYX) is a new standard for exchanging deep learning models. It promises to make deep learning models portable thus preventing vendor lock in. Let’s look at why that matters for the modern ML/AI developer.

The end result of a trained deep learning algorithm is a model file that efficiently represents the relationship between input data and output predictions. A neural network is one of the most powerful ways to generate these predictive models but can be difficult to build in to production systems. Most often, these models exist in a data format such as a `.pth` file or an HD5 file. Oftentimes you want these models to be portable so that you can deploy them in environments that might be different than where you initially trained the model.

ONNX Overview

At a high level, ONNX is designed to allow framework interoporability. There are many excellent machine learning libraries in various languages — PyTorch, TensorFlow, MXNet, and Caffe are just a few that have become very popular in recent years, but there are many others as well.
The idea is that you can train a model with one tool stack and then deploy it using another for inference and prediction. To ensure this interoperability you must export your model in the model.onnx format which is serialized representation of the model in a protobuf file. Currently there is native support in ONNX for PyTorch, CNTK, MXNet, and Caffe2 but there are also converters for TensorFlow and CoreML.

Prerequisites

• Basic understanding of machine learning (ML) and deep learning (DL): Familiarity with concepts like neural networks, layers, models, and training processes.

• Knowledge of popular ML frameworks: Experience with at least one major ML/DL framework like TensorFlow, PyTorch, or Scikit-learn, since ONNX is used to convert models between frameworks.

• Model deployment experience: Understanding of the basics of deploying ML models in production environments, as ONNX aids in efficient cross-platform deployment.

• Familiarity with model optimization: Basic understanding of techniques to optimize ML models for performance, as ONNX is often used to make models more efficient.

• Python programming: Since ONNX has strong support in Python, knowledge of Python is essential for using its tools and libraries.

ONNX in Practice

Let’s imagine that you want to train a model to predict if a food item in your refrigerator is still good to eat. You decide to run a a bunch of photos of food that is at various stages past its expiration date and pass it in to a convolutional neural network (CNN) that looks at images of food and trains it to predict if the food is still edible.

Once you have trained your model, you then want to deploy it to a new iOS app so that anyone can use your pre-trained model to check their own food for safety. You initially trained your model using PyTorch but iOS expects to use CoreML to be used inside the app. ONNX is an intermediary representation of your model that lets you easily go from one environment to the next.

Using PyTorch you would normally export your model using torch.save(the_model.state_dict(), PATH) Exporting to the ONNX interchange format is just one more line:
torch.onnx.export(model, dummy_input, 'SplitModel.proto', verbose=True)

Using a tool like ONNX-CoreML, you can now easily turn your pre-trained model in to a file that you can import in to XCode and integrate seamlessly with your app. For a working example, checkout this excellent post by Stefano Attardi on building a ML-driven iOS app from start to finish.

Conclusion

As more and more deep learning frameworks emerge and workflows become more advanced, the need for portability is more important than ever. ONNX is a powerful and open standard for preventing framework lock-in and ensuring that you the models you develop will be usable in the long run.