Neural Network represented by ONNX

Introduction
ONNXRuntime
Profiling the Model
Other Tools
- Netron
- ONNX GraphSurgeon

Introduction

A machine learning model is often represetned as a computation graph in the ONNX format.¹

Operationally, we train our model and then (somehow) export the model to a *.onnx file on our harddrive. This file includes the definition of our model, a computation graph², and the corresponding weights.

The building blocks of a ONNX computation graph are:

Nodes: the input/output tensors or the operators (a single unit to carry out a specific calculation).
Initializers: the constants³ within the graph, often act as the weight of the model.
Attributes: Fixed parameters of an operator⁴.

For instance, a linear regression model: y = a x + c, could be represented with,

inputs
- x: a input tensor
- a: a initializer
- c: a initializer
output: tensor y
operators:
- MatMul: take two inputs, x and a
- Add: take two inputs, the output of MatMul and c

ONNXRuntime

The onnxruntime library provides C++/Python/C# APIs to execute an ONNX model, so that the inference process can be executed on various platforms including CPU and GPU.

The onnxruntime library also provides some model optimization methods like basic operator fusion⁵ and quantization⁶.

The very nice feature about onnxruntime is its support of various execution providers (EP). Effectively this means we can write a single standard codebase with the onnxruntime API, and make it run on different devices, simply by changing the EP in the configuration. Sounds good, right?

As far as I know, the standard release version⁷ only contains 3 EPs:

CPUExecutionProvider: support the model to be executed on a general CPU⁸.
CUDAExecutionProvider: support the model to be executed on a Nvidia GPU.
TensorrtExecutionProvider: support the model to be executed on a Nvidia GPU, but in a (significantly) optimized engine named TensorRT.⁹

If we want to use more EPs, we have to build our own onnxruntime from the source code. For instance, if we want to use the OpenVINOExecutionProvider, we have to perform two steps,

Download the source code of openvino and build it.
Download the source code of onnxruntime and build it, following the EP-specific building instruction.

Here is an example of step 2, where I used the code to build onnxruntime with OpenVINOExecutionProvider supported.

DIR_OPENVINO=../openvino
DIR_CUDA=$HOME/.local/cuda
DIR_TRT=$HOME/.local/tensorrt

source $DIR_OPENVINO/setupvar.sh
./build.sh --config Release \
          --build_shared_lib \
          --parallel \
          --skip_submodule_sync \
          --use_cuda \
          --cuda_home $DIR_CUDA \
          --cudnn_home $DIR_CUDA \
          --use_tensorrt --tensorrt_home $DIR_TRT \
          --use_openvino CPU_FP32 \
          --build_wheel

We can check the availabe EPs with the Python API provided by onnxruntime

In [1]: import onnxruntime

In [2]: onnxruntime.get_available_providers()
Out[2]: ['OpenVINOExecutionProvider', 'CPUExecutionProvider']

Profiling the Model

We want our neural network code to execute fast. We do this by working hard on the slow part. We know which part of our model is slow by profiling the execution/inference process.

If we use onnxruntime to execute our ONNX model, we can enable profiling by setting the attribute enable_profiling of the SessionOptions instance to True. Here is an exmple.

import onnxruntime as ort

sess_options = ort.SessionOptions()
sess_options.enable_profiling = True

session = ort.InferenceSession(
    "your_model.onnx",
    sess_options=sess_options,
)

# Run your model ...

session.end_profiling()

The code above would generate a JSON file (.json), which recorded the start time and end time for each nodes.

It is convenient to use the Chrome tracing tool to visualise the json file. To start the tool, we type the following URL on a Chrome/Chromium browser,

chrome://tracing

And load the JSON file by clicking the load button.

Other Tools

Netron

A viewer for neural network models, including ONNX.

You can just open the .onnx file with Netron and see the computation graph rendered as, well, a graph!

It is fun (and challenging) to follow the computation graph in the rendered picture in Netron. I encourage everyone to trace the input tensor all the way through the graph once. You are likely to learn something new!

ONNX GraphSurgeon

Nvidia’s Onnx GraphSurgeon is a Python package for editing the computation graph represented by ONNX.

It is basically a wrapper around the onnx package, but the APIs are more intuitive, well, at least for me.

It is distributed with the TensorRT package, but it is self contained. You can install it to edit your .onnx file without installing TensorRT.

we will use the term “model”, “computation graph”, “onnx model”, and “neural network” interchangably in this note. ↩
Mathematically, a graph is composed of nodes and edges. For machine learning models, the computation graph corresponds to the execution instructions to compute the required result. ↩
the name intializer is unfortunately misleading. In fact it is called constant in onnx_graphsurgeon package developed by Nvidia, which is a wrapper of the onnx package. ↩
for instance, the Gemm operator has parameters alpha, beta, transA, and transB. ↩
As far as I know, onnxruntime automatically fuse MatMul and Add into GEMM. ↩
In the default setting, onnxruntime insert quantization and de-quantization operators for every node in the computation graph, to reduce the precision of the model weight from 32-bit float point (float in C++, np.float32 in Python) to 8-bit signed integer. ↩
You can get the release easily from PyPI with pip install onnxruntime-gpu ↩
I only used onnxruntime on the X68 architecture. I never tested it on ARM or RISC-V. I suspect the onnxruntime folks would support ARM given the success of Macbooks. But I am not privileged to get a good Mac for myself! ↩
TensorRT is extremely powerful. I have witnessed 5x to 10x performance boost simply by converting a .onnx file to a TensorRT engine, with a reduced GPU memory consumption. ↩