Tutorial

What every ML/AI developer should know about ONNX

Updated on September 25, 2024
author

Dillon

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.

Thanks for learning with the DigitalOcean Community. Check out our offerings for compute, storage, networking, and managed databases.

Learn more about our products

About the authors
Default avatar
Dillon

author

Still looking for an answer?

Ask a questionSearch for more help

Was this helpful?
 
Leave a comment


This textbox defaults to using Markdown to format your answer.

You can type !ref in this text area to quickly search our full set of tutorials, documentation & marketplace offerings and insert the link!

Try DigitalOcean for free

Click below to sign up and get $200 of credit to try our products over 60 days!

Sign up

Join the Tech Talk
Success! Thank you! Please check your email for further details.

Please complete your information!

Become a contributor for community

Get paid to write technical tutorials and select a tech-focused charity to receive a matching donation.

DigitalOcean Documentation

Full documentation for every DigitalOcean product.

Resources for startups and SMBs

The Wave has everything you need to know about building a business, from raising funding to marketing your product.

Get our newsletter

Stay up to date by signing up for DigitalOcean’s Infrastructure as a Newsletter.

New accounts only. By submitting your email you agree to our Privacy Policy

The developer cloud

Scale up as you grow — whether you're running one virtual machine or ten thousand.

Get started for free

Sign up and get $200 in credit for your first 60 days with DigitalOcean.*

*This promotional offer applies to new accounts only.