Creating and Modifying ONNX Model Using ONNX Python API
Introduction
Open Neural Network Exchange (ONNX) is an open standard format for representing machine learning models. ONNX is the most widely used machine learning model format, supported by a community of partners who have implemented it in many frameworks and tools.
In this blog post, I would like to discuss how to use the ONNX Python API to create and modify ONNX models.
ONNX Data Structure
ONNX model is represented using protocol buffers. Specifically, the entire model information was encoded using the onnx.proto.
The major ONNX protocol buffers to describe a neural network are ModelProto, GraphProto, NodeProto, TensorProto, ValueInfoProto.
Key ONNX Protos | Description |
---|---|
ModelProto | It contains model description and GraphProto. |
GraphProto | It contains the node information, node initializers, and IO tensors in the model. |
NodeProto | It represents a node in the model. It contains the input and output tensor names, node initializers, and node attributes. |
TensorProto | It represents an node initializer (constant tensor in the node). In addition to the data type and shape, specific values were assigned. |
ValueInfoProto | It represents an IO tensor in the model in which only the data type and shape were defined. |
Creating ONNX Model
To better understand the ONNX protocol buffers, let’s create a dummy convolutional classification neural network, consisting of convolution, batch normalization, ReLU, average pooling layers, from scratch using ONNX Python API (ONNX helper functions onnx.helper).
1 | import numpy as np |
Once the ONNX model is created, we can further verify the model using ONNX Runtime.
Modifying ONNX Model
Modifying ONNX model is a little bit complicated since all the information were encoded as protocol buffers and there is no ONNX helper function to modify the protocol buffers.
Fortunately, we can assign values to the non-repeated
attributes in the onnx.proto directly. For the repeated
attributes, we cannot assign new values to it, but we are allowed to modify the values in place, or use Python binding interface, such as extend
and pop
, to add and remove items.
1 | import onnx |
Source Code
The source code of the implementation is available on GitHub.
Miscellaneous
Dealing with ONNX protocol buffer is complicated and error-prone. The ONNX protocol buffer representation also depends on ONNX IR version and opset version. It will be desirable in some scenarios if we can have a high-level abstracted interface that allows us to modify the model without having to going through the low-level data structure. Fortunately, we have ONNX GraphSurgeon that can help us to do so.
References
Creating and Modifying ONNX Model Using ONNX Python API