What every ML/AI developer should know about ONNX

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.

6 years ago   •   2 min read

By Dillon

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.

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.

Spread the word

Keep reading