第十一章 ONNX 使用
第十一章簡介
本章介紹模型部署的第一個工具——ONNX。
ONNX （Open Neural Network Exchange，開放神經網路交換格式）是一種開放的、跨平臺的深度學習模型交換格式，可以方便地將模型從一個框架轉移到另一個框架，可謂是模型部署必須要瞭解的一個工具。
本章將從ONNX的概念及原理介紹開始，再介紹ONNX配套的推理引擎——ONNXRuntime， 最後介紹ONNXRuntime中常用的優化方法（float16量化、int8量化、混合精度量化、計算圖優化、執行緒管理和IO binding）。
11.1 ONNX 簡介與安裝
前言
在深度學習演算法開發過程中，模型訓練與部署是兩個環節，pytorch通常只用於訓練，獲得模型權重檔，而最終部署還有專門的部署平臺，例如TensorRT、NCNN、OpenVINO等幾十種部署推理平臺。
如何將pytorch模型檔讓幾十種部署推理平臺能接收與讀取是個大問題。即使各推理平臺都適配pytorch，那還有其他訓練框架也要適配，是非常麻煩的。
假設有N個訓練框架，M個推理框架，互相要適配，那就是O(NM)的複雜度。如果能有一種中間格式作為一個標注，能被所有框架所適配，那複雜度順便降低為O(N+M)。
onnx就是為了降低深度學習模型從訓練到部署的複雜度，由微軟和meta在2017年提出的一種開放神經網路交換格式，目的在於方便的將模型從一個框架轉移到另一個框架。
本小結就介紹onnx基礎概念、pytorch模型匯出onnx模型。
ONNX 簡介
ONNX （Open Neural Network Exchange，開放神經網路交換格式）是一種開放的、跨平臺的深度學習模型交換格式，可以方便地將模型從一個框架轉移到另一個框架。
onnx最初由微軟和meta在2017年聯合發佈，後來亞馬遜也加入進來，目前已經成為行業共識，目前已經有50多個機構的產品支援onnx。
onnx最大的優點是簡化了模型部署之間因框架的不同帶來的繁瑣事，這就像普通話。在中國129種方言之間要互相通信是很困難的，解決辦法就是設計一種可以與129種語言進行轉換的語言——普通話。onnx就是一個支援絕大多數主流機器學習模型格式之間轉換的格式。
採用pytorch進行模型開發時，部署環節通常將pytorch模型轉換為onnx模型，然後再進行其他格式轉換，或者直接採用onnx檔進行推理，在本章節就介紹採用onnx檔進行推理的方法。
ONNX 基礎概念
onnx檔是一種計算圖，用於描述資料要進行何種計算，它就像是數學計算的語言，可以進行計算的操作稱之為操作符——operator，一系列operator構成一個計算圖。
計算圖中包含了各節點、輸入、輸出、屬性的詳細資訊，有助於開發者觀察模型結構。
下面通過一個線性回歸模型的計算圖來瞭解onnx的計算圖
可以採用python代碼構建onnx計算圖，運行配套代碼，構建了一個線性回歸模型
from onnx import TensorProto
from onnx.helper import (
make_model, make_node, make_graph,
make_tensor_value_info)
# 'X' is the name, TensorProto.FLOAT the type, [None, None] the shape
X = make_tensor_value_info('X', TensorProto.FLOAT, [None, None])
A = make_tensor_value_info('A', TensorProto.FLOAT, [None, None])
B = make_tensor_value_info('B', TensorProto.FLOAT, [None, None])
Y = make_tensor_value_info('Y', TensorProto.FLOAT, [None])
node1 = make_node('MatMul', ['X', 'A'], ['XA'])
node2 = make_node('Add', ['XA', 'B'], ['Y'])
graph = make_graph([node1, node2], # nodes
'lr', # a name
[X, A, B], # inputs
[Y]) # outputs
onnx_model = make_model(graph)
with open("linear_regression.onnx", "wb") as f:
f.write(onnx_model.SerializeToString())
Copy
運行以上代碼會獲得linear_regression.onnx檔，可通過https://netron.app/ 進行視覺化
圖中
- A, B, X, Y表示輸入輸出資料
- 黑色的MatMul和Add是Node，表示具體的操作
- format：表示生成該onnx檔的onnx版本
- imports：operator的版本；運算元是onnx中最重要的一個概念，大多數模型不成功是因為沒有對應的運算元，因此運算元集的版本選擇很重要；
- inputs和outputs：是輸入和輸出，其中type是資料類型以及shape。
為了進一步瞭解onnx檔，下面匯出一個resnet50進行觀察，onnx檔可通過以下代碼獲得：
import torchvision
import torch
model = torchvision.models.resnet50(pretrained=False)
dummy_data = torch.randn((1, 3, 224, 224))
with torch.no_grad():
torch.onnx.export(model, (dummy_data),
"resnet50.onnx",
opset_version=11,
input_names=['input_name_edit_by_tingsongyu'],
output_names=['output_name_edit_by_tingsongyu'])
Copy
下面再看一個resnet50的onnx檔，觀察更多運算元的描述。
更多onnx基礎概念參見官網：https://onnx.ai/onnx/intro/concepts.html
ONNX 的operator
上面介紹了onnx檔主要定義了計算圖，計算圖中的每個操作稱為運算元，運算元庫的豐富程度，直接決定了onnx可以表示模型的種類。
關於onnx支持哪些運算元，一定要上官網看一看。
對於普通使用者，需要關注使用時的opset是哪個版本，目前最新版本是20。運算元庫可通過以下函數查看。
import onnx
print(onnx.version, " opset=", onnx.defs.onnx_opset_version())
Copy
關於運算元的理解，以及不適配問題，推薦OpenMMLab的三篇博文
https://zhuanlan.zhihu.com/p/479290520：講解了pytorch轉onnx時，每一個操作是如何轉換到onnx運算元的；介紹了運算元映射關係
https://zhuanlan.zhihu.com/p/513387413：講解了三種添加運算元的方法
其中有一張圖對於理解pytorch轉onnx很有幫助，這裡引用一下：
ONNX 安裝
onnx的安裝很簡單：pip install onnx
在這裡，提前說一下，onnx是onnx，與onnxruntime不是同一個東西，它們要分開安裝，也要分開理解。
pytorch匯出onnx
pytorch模型匯出為onnx調用torch.onnx.export函數即可，該函數包含很多參數，這裡只介紹幾個常用的，更多的參考官方文檔
torch.onnx.export(model, args, f, export_params=True, verbose=False, training=, input_names=None, output_names=None, operator_export_type=, opset_version=None, do_constant_folding=True, dynamic_axes=None, keep_initializers_as_inputs=None, custom_opsets=None, export_modules_as_functions=False)
- model: 需要被轉換的模型，可以有三種類型， torch.nn.Module, torch.jit.ScriptModule or torch.jit.ScriptFunction
- args：model輸入時所需要的參數，這裡要傳參時因為構建計算圖過程中，需要採用資料對模型進行一遍推理，然後記錄推理過程需要的操作，然後生成計算圖。args要求是tuple或者是Tensor的形式。一般只有一個輸入時，直接傳入Tensor，多個輸入時要用tuple包起來。
- export_params: 是否需要保存參數。預設為True，通常用於模型結構遷移到其它框架時，可以用False。
- input_names：輸入資料的名字， (list of str, default empty list) ，在使用onnx檔時，資料的傳輸和使用，都是通過name: value的形式。
- output_names：同上。
- opset_version：使用的運算元集版本。
- dynamic_axes：動態維度的指定，例如batchsize在使用時隨時會變，則需要把該維度指定為動態的。預設情況下計算圖的資料維度是固定的，這有利於效率提升，但缺乏靈活性。用法是，對於動態維度的輸入、輸出，需要設置它哪個軸是動態的，並且為這個軸設定名稱。這裡有3個要素，資料名稱，軸序號，軸名稱。因此是通過dict來設置的。例如dynamic_axes={ "x": {0: "my_custom_axis_name"} }， 表示名稱為x的資料，第0個軸是動態的，動態軸的名字叫my_custom_axis_name。通常用於batchsize或者是對於h,w是不固定的模型要設置動態軸。
接下來以resnet50為例，匯出一個在ImageNet上訓練好的分類模型，再通過netron觀察區別。
下面使用配套代碼匯出三個模型，分別是bs=1, bs=128， bs為動態的，下一節將對比兩者效率。
import torchvision
import torch
model = torchvision.models.resnet50(weights=torchvision.models.ResNet50_Weights.IMAGENET1K_V1)
if __name__ == '__main__':
op_set = 13
dummy_data = torch.randn((1, 3, 224, 224))
dummdy_data_128 = torch.randn((128, 3, 224, 224))
# 固定 batch = 1
torch.onnx.export(model, (dummy_data), "resnet50_bs_1.onnx",
opset_version=op_set, input_names=['input'], output_names=['output'])
# 固定 batch = 128
torch.onnx.export(model, (dummdy_data_128), "resnet50_bs_128.onnx",
opset_version=op_set, input_names=['input'], output_names=['output'])
# 動態 batch
torch.onnx.export(model, (dummy_data), "resnet50_bs_dynamic.onnx",
opset_version=op_set, input_names=['input'], output_names=['output'],
dynamic_axes={"input": {0: "batch_axes"},
"output": {0: "batch_axes"}})
Copy
對比如下圖所示，input的type中，shape一個是1，一個是batch_axes，其中batch_axes這個就是自訂的命名。
小結
本小節介紹了onnx提出的目的與意義，還有基礎概念，onnx可以作為一個中間格式，被絕大多數框架所適配，方便開發人員從訓練框架轉到開發框架。
onnx檔核心是記錄模型的計算圖，包括輸入資料、各操作節點、輸出資料等資訊。
最後介紹了pytorch匯出onnx的方法，其中需要主要的是op_set版本，以及動態維度的設置。
下一小節，將利用本節匯出的onnx模型檔，在onnx的推理庫——onnxruntime上進行推理以及性能效率評估。
留言列表
