文章首發於我的 github 倉庫-cv演算法工程師成長之路，歡迎關注我的公眾號-嵌入式視覺。
本文大部分內容為對 ONNX 官方資料的總結和翻譯，部分知識點參考網上質量高的部落格。
一，ONNX 概述
深度學習演算法大多透過計算資料流圖來完成神經網路的深度學習過程。 一些框架（例如CNTK，Caffe2，Theano和TensorFlow）使用靜態圖形，而其他框架（例如 PyTorch 和 Chainer）使用動態圖形。 但是這些框架都提供了介面，使開發人員可以輕鬆構建計算圖和執行時，以最佳化的方式處理圖。 這些圖用作中間表示（IR），捕獲開發人員原始碼的特定意圖，有助於最佳化和轉換在特定裝置（CPU，GPU，FPGA等）上執行。
ONNX 的本質只是一套開放的 ML 模型標準，模型檔案儲存的只是網路的拓撲結構和權重（其實每個深度學習框架最後儲存的模型都是類似的），脫離開框架是沒辦法對模型直接進行 inference的。
1.1，為什麼使用通用 IR
現在很多的深度學習框架提供的功能都是類似的，但是在 API、計算圖和 runtime 方面卻是獨立的，這就給 AI 開發者在不同平臺部署不同模型帶來了很多困難和挑戰，ONNX 的目的在於提供一個跨框架的模型中間表達框架，用於模型轉換和部署。ONNX 提供的計算圖是通用的，格式也是開源的。
二，ONNX 規範
Open Neural Network Exchange Intermediate Representation (ONNX IR) Specification.
ONNX 結構的定義檔案 .proto 和 .prpto3 可以在 onnx folder 目錄下找到，檔案遵循的是谷歌 Protobuf 協議。ONNX 是一個開放式規範，由以下元件組成：
- 可擴充套件計算圖模型的定義
- 標準資料型別的定義
- 內建運算子的定義
IR6 版本的 ONNX 只能用於推理（inference），從 IR7 開始 ONNX 支援訓練（training）。onnx.proto 主要的物件如下：
- ModelProto
- GraphProto
- NodeProto
- AttributeProto
- ValueInfoProto
- TensorProto
他們之間的關係：ONNX 模型 load 之後，得到的是一個 ModelProto，它包含了一些版本資訊，生產者資訊和一個非常重要的 GraphProto；在 GraphProto 中包含了四個關鍵的 repeated 陣列，分別是node (NodeProto 型別)，input(ValueInfoProto 型別)，output(ValueInfoProto 型別)和 initializer (TensorProto 型別)，其中 node 中存放著模型中的所有計算節點，input 中存放著模型所有的輸入節點，output 存放著模型所有的輸出節點，initializer 存放著模型所有的權重；節點與節點之間的拓撲定義可以透過 input 和output 這兩個 string 陣列的指向關係得到，這樣利用上述資訊我們可以快速構建出一個深度學習模型的拓撲圖。最後每個計算節點當中還包含了一個 AttributeProto 陣列，用於描述該節點的屬性，例如 Conv 層的屬性包含 group，pads 和strides 等等，具體每個計算節點的屬性、輸入和輸出可以參考這個 Operators.md 文件。
需要注意的是，上面所說的 GraphProto 中的 input 輸入陣列不僅僅包含我們一般理解中的圖片輸入的那個節點，還包含了模型當中所有權重。舉例，Conv 層中的 W 權重實體是儲存在 initializer 當中的，那麼相應的會有一個同名的輸入在 input 當中，其背後的邏輯應該是把權重也看作是模型的輸入，並透過 initializer 中的權重實體來對這個輸入做初始化(也就是把值填充進來)
2.1，Model
模型結構的主要目的是將後設資料( meta data)與圖形(graph)相關聯，圖形包含所有可執行元素。 首先，讀取模型檔案時使用後設資料，為實現提供所需的資訊，以確定它是否能夠：執行模型，生成日誌訊息，錯誤報告等功能。此外後設資料對工具很有用，例如IDE和模型庫，它需要它來告知使用者給定模型的目的和特徵。
每個 model 有以下元件：
| Name | Type | Description |
|---|---|---|
| ir_version | int64 | The ONNX version assumed by the model. |
| opset_import | OperatorSetId | A collection of operator set identifiers made available to the model. An implementation must support all operators in the set or reject the model. |
| producer_name | string | The name of the tool used to generate the model. |
| producer_version | string | The version of the generating tool. |
| domain | string | A reverse-DNS name to indicate the model namespace or domain, for example, 'org.onnx' |
| model_version | int64 | The version of the model itself, encoded in an integer. |
| doc_string | string | Human-readable documentation for this model. Markdown is allowed. |
| graph | Graph | The parameterized graph that is evaluated to execute the model. |
| metadata_props | map<string,string> | Named metadata values; keys should be distinct. |
| training_info | TrainingInfoProto[] | An optional extension that contains information for training. |
2.2，Operators Sets
每個模型必須明確命名它依賴於其功能的運算子集。 操作員集定義可用的運算子，其版本和狀態。 每個模型按其域定義匯入的運算子集。 所有模型都隱式匯入預設的 ONNX 運算子集。
運算子集(Operators Sets)物件的屬性如下：
| Name | Type | Description |
|---|---|---|
| magic | string | T ‘ONNXOPSET’ |
| ir_version | int32 | The ONNX version corresponding to the operators. |
| ir_version_prerelease | string | The prerelease component of the SemVer of the IR. |
| ir_build_metadata | string | The build metadata of this version of the operator set. |
| domain | string | The domain of the operator set. Must be unique among all sets. |
| opset_version | int64 | The version of the operator set. |
| doc_string | string | Human-readable documentation for this operator set. Markdown is allowed. |
| operator | Operator[] | The operators contained in this operator set. |
2.3，ONNX Operator
圖( graph)中使用的每個運算子必須由模型(model)匯入的一個運算子集明確宣告。
運算子（Operator）物件定義的屬性如下：
| Name | Type | Description |
|---|---|---|
| op_type | string | The name of the operator, as used in graph nodes. MUST be unique within the operator set’s domain. |
| since_version | int64 | The version of the operator set when this operator was introduced. |
| status | OperatorStatus | One of ‘EXPERIMENTAL’ or ‘STABLE.’ |
| doc_string | string | A human-readable documentation string for this operator. Markdown is allowed. |
2.4，ONNX Graph
序列化圖由一組後設資料欄位(metadata)，模型引數列表(a list of model parameters,)和計算節點列表組成(a list of computation nodes)。每個計算資料流圖被構造為拓撲排序的節點列表，這些節點形成圖形，其必須沒有周期。 每個節點代表對運營商的呼叫。 每個節點具有零個或多個輸入以及一個或多個輸出。
圖表(Graph)物件具有以下屬性：
| Name | Type | Description |
|---|---|---|
| name | string | 模型計算圖的名稱 |
| node | Node[] | 節點列表，基於輸入/輸出資料依存關係形成部分排序的計算圖，拓撲順序排列。 |
| initializer | Tensor[] | 命名張量值的列表。 當 initializer 與計算圖 graph輸入名稱相同，輸入指定一個預設值，否則指定一個常量值。 |
| doc_string | string | 用於閱讀模型的文件 |
| input | ValueInfo[] | 計算圖 graph 的輸入引數，在 ‘initializer.’ 中可能能找到預設的初始化值。 |
| output | ValueInfo[] | 計算圖 graph 的輸出引數。 |
| value_info | ValueInfo[] | 用於儲存除輸入、輸出值之外的型別和形狀資訊。 |
2.5，ValueInfo
ValueInfo 物件屬性如下：
| Name | Type | Description |
|---|---|---|
| name | string | The name of the value/parameter. |
| type | Type | The type of the value including shape information. |
| doc_string | string | Human-readable documentation for this value. Markdown is allowed. |
2.6，Standard data types
ONNX 標準有兩個版本，主要區別在於支援的資料型別和運算元不同。計算圖 graphs、節點 nodes和計算圖的 initializers 支援的資料型別如下。原始數字，字串和布林型別必須用作張量的元素。
2.6.1，Tensor Element Types
| Group | Types | Description |
|---|---|---|
| Floating Point Types | float16, float32, float64 | 浮點數遵循IEEE 754-2008標準。 |
| Signed Integer Types | int8, int16, int32, int64 | 支援 8-64 位寬的有符號整數。 |
| Unsigned Integer Types | uint8, uint16 | 支援 8 或 16 位的無符號整數。 |
| Complex Types | complex64, complex128 | 具有 32 位或 64 位實部和虛部的複數。 |
| Other | string | 字串代表的文字資料。 所有字串均使用UTF-8編碼。 |
| Other | bool | 布林值型別，表示的資料只有兩個值，通常為 true 和 false。 |
2.6.2，Input / Output Data Types
以下型別用於定義計算圖和節點輸入和輸出的型別。
| Variant | Type | Description |
|---|---|---|
| ONNX | dense tensors | 張量是向量和矩陣的一般化 |
| ONNX | sequence | sequence (序列)是有序的稠密元素集合。 |
| ONNX | map | 對映是關聯表，由鍵型別和值型別定義。 |
ONNX 現階段沒有定義稀疏張量型別。
三，ONNX版本控制
四，主要運算元概述
五，Python API 使用
5.1，載入模型
1，Loading an ONNX model
import onnx
# onnx_model is an in-mempry ModelProto
onnx_model = onnx.load('path/to/the/model.onnx') # 載入 onnx 模型
2，Loading an ONNX Model with External Data
- 【預設載入模型方式】如果外部資料(
external data)和模型檔案在同一個目錄下，僅使用onnx.load()即可載入模型，方法見上小節。 - 如果外部資料(
external data)和模型檔案不在同一個目錄下，在使用onnx_load()函式後還需使用load_external_data_for_model()函式指定外部資料路徑。
import onnx
from onnx.external_data_helper import load_external_data_for_model
onnx_model = onnx.load('path/to/the/model.onnx', load_external_data=False)
load_external_data_for_model(onnx_model, 'data/directory/path/')
# Then the onnx_model has loaded the external data from the specific directory
3，Converting an ONNX Model to External Data
from onnx.external_data_helper import convert_model_to_external_data
# onnx_model is an in-memory ModelProto
onnx_model = ...
convert_model_to_external_data(onnx_model, all_tensors_to_one_file=True, location='filename', size_threshold=1024, convert_attribute=False)
# Then the onnx_model has converted raw data as external data
# Must be followed by save
5.2，儲存模型
1，Saving an ONNX Model
import onnx
# onnx_model is an in-memory ModelProto
onnx_model = ...
# Save the ONNX model
onnx.save(onnx_model, 'path/to/the/model.onnx')
2，Converting and Saving an ONNX Model to External Data
import onnx
# onnx_model is an in-memory ModelProto
onnx_model = ...
onnx.save_model(onnx_model, 'path/to/save/the/model.onnx', save_as_external_data=True, all_tensors_to_one_file=True, location='filename', size_threshold=1024, convert_attribute=False)
# Then the onnx_model has converted raw data as external data and saved to specific directory
5.3，Manipulating TensorProto and Numpy Array
import numpy
import onnx
from onnx import numpy_helper
# Preprocessing: create a Numpy array
numpy_array = numpy.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=float)
print('Original Numpy array:\n{}\n'.format(numpy_array))
# Convert the Numpy array to a TensorProto
tensor = numpy_helper.from_array(numpy_array)
print('TensorProto:\n{}'.format(tensor))
# Convert the TensorProto to a Numpy array
new_array = numpy_helper.to_array(tensor)
print('After round trip, Numpy array:\n{}\n'.format(new_array))
# Save the TensorProto
with open('tensor.pb', 'wb') as f:
f.write(tensor.SerializeToString())
# Load a TensorProto
new_tensor = onnx.TensorProto()
with open('tensor.pb', 'rb') as f:
new_tensor.ParseFromString(f.read())
print('After saving and loading, new TensorProto:\n{}'.format(new_tensor))
5.4，建立ONNX模型
可以透過 helper 模組提供的函式 helper.make_graph 完成建立 ONNX 格式的模型。建立 graph 之前，需要先建立相應的 NodeProto(node)，參照文件設定節點的屬性，指定該節點的輸入與輸出，如果該節點帶有權重那還需要建立相應的ValueInfoProto 和 TensorProto 分別放入 graph 中的 input 和 initializer 中，以上步驟缺一不可。
import onnx
from onnx import helper
from onnx import AttributeProto, TensorProto, GraphProto
# The protobuf definition can be found here:
# https://github.com/onnx/onnx/blob/master/onnx/onnx.proto
# Create one input (ValueInfoProto)
X = helper.make_tensor_value_info('X', TensorProto.FLOAT, [3, 2])
pads = helper.make_tensor_value_info('pads', TensorProto.FLOAT, [1, 4])
value = helper.make_tensor_value_info('value', AttributeProto.FLOAT, [1])
# Create one output (ValueInfoProto)
Y = helper.make_tensor_value_info('Y', TensorProto.FLOAT, [3, 4])
# Create a node (NodeProto) - This is based on Pad-11
node_def = helper.make_node(
'Pad', # name
['X', 'pads', 'value'], # inputs
['Y'], # outputs
mode='constant', # attributes
)
# Create the graph (GraphProto)
graph_def = helper.make_graph(
[node_def], # nodes
'test-model', # name
[X, pads, value], # inputs
[Y], # outputs
)
# Create the model (ModelProto)
model_def = helper.make_model(graph_def, producer_name='onnx-example')
print('The model is:\n{}'.format(model_def))
onnx.checker.check_model(model_def)
print('The model is checked!')
5.5，檢查模型
在完成 ONNX 模型載入或者建立後，有必要對模型進行檢查，使用 onnx.check.check_model() 函式。
import onnx
# Preprocessing: load the ONNX model
model_path = 'path/to/the/model.onnx'
onnx_model = onnx.load(model_path)
print('The model is:\n{}'.format(onnx_model))
# Check the model
try:
onnx.checker.check_model(onnx_model)
except onnx.checker.ValidationError as e:
print('The model is invalid: %s' % e)
else:
print('The model is valid!')
5.6，實用功能函式
函式 extract_model() 可以從 ONNX 模型中提取子模型，子模型由輸入和輸出張量的名稱定義。這個功能方便我們 debug 原模型和轉換後的 ONNX 模型輸出結果是否一致(誤差小於某個閾值)，不再需要我們手動去修改 ONNX 模型。
import onnx
input_path = 'path/to/the/original/model.onnx'
output_path = 'path/to/save/the/extracted/model.onnx'
input_names = ['input_0', 'input_1', 'input_2']
output_names = ['output_0', 'output_1']
onnx.utils.extract_model(input_path, output_path, input_names, output_names)
5.7，工具
函式 update_inputs_outputs_dims() 可以將模型輸入和輸出的維度更新為引數中指定的值，可以使用 dim_param 提供靜態和動態尺寸大小。
import onnx
from onnx.tools import update_model_dims
model = onnx.load('path/to/the/model.onnx')
# Here both 'seq', 'batch' and -1 are dynamic using dim_param.
variable_length_model = update_model_dims.update_inputs_outputs_dims(model, {'input_name': ['seq', 'batch', 3, -1]}, {'output_name': ['seq', 'batch', 1, -1]})
# need to check model after the input/output sizes are updated
onnx.checker.check_model(variable_length_model )
參考資料
文章同步發於 github、知乎，最新版以github為主。
本人水平有限，文章如有問題，歡迎及時指出。如果看完文章有所收穫，一定要先點贊後收藏。畢竟，贈人玫瑰，手有餘香。
最後，更多面經和乾貨文章，微信搜尋我的公眾號-嵌入式視覺!