What Every ML/AI Developer Should Know About ONNX

Introduction

The Open Neural Network Exchange Format (ONNX) 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 into 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 from where you initially trained the model.

ONNX Overview

At a high level, ONNX is designed to allow framework interoperability. 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 a 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 for ML/AI Developer

• 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. 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 for ML/AI Developer

Imagine you want to train a model to predict if a food item in your refrigerator is still edible. You use photos of food at various expiration stages and input them into a convolutional neural network. This network analyzes the images and trains the model to predict whether the food is safe to eat.

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 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)
    

With ONNX-CoreML, you can easily convert your pre-trained model into a file for seamless XCode integration. Check out Stefano Attardi's excellent post on building a complete ML-driven iOS app from start to finish.

Conclusion for What Every ML/AI Developer Should Know About ONNX

As deep learning frameworks grow and workflows advance, portability becomes increasingly essential. ONNX is a powerful and open standard for preventing framework lock-in. It ensures that the models you develop remain usable in the long run.