pytorch模型轉onnx模型的方法詳解

2022-08-30 18:04:01

學習目標

1.掌握pytorch模型轉換到onnx模型

2.順利執行onnx模型

3.比對onnx模型和pytorch模型的輸出結果

學習大綱

pytorch模型轉換onnx模型
執行onnx模型
onnx模型輸出與pytorch模型比對

學習內容

前提條件：需要安裝onnx 和 onnxruntime，可以通過 pip install onnx 和 pip install onnxruntime 進行安裝

1 . pytorch 轉 onnx

pytorch 轉 onnx 只需要一個函數 torch.onnx.export

torch.onnx.export(model, args, path, export_params, verbose, input_names, output_names, do_constant_folding, dynamic_axes, opset_version)

引數說明：

model——需要匯出的pytorch模型
args——模型的輸入引數，滿足輸入層的shape正確即可。
path——輸出的onnx模型的位置。例如‘yolov5.onnx’。
export_params——輸出模型是否可訓練。default=True，表示匯出trained model，否則untrained。
verbose——是否列印模型轉換資訊。default=False。
input_names——輸入節點名稱。default=None。
output_names——輸出節點名稱。default=None。
do_constant_folding——是否使用常數摺疊（不瞭解），預設即可。default=True。
dynamic_axes——模型的輸入輸出有時是可變的，如Rnn，或者輸出影象的batch可變，可通過該引數設定。如輸入層的shape為（b，3，h，w），batch，height，width是可變的，但是chancel是固定三通道。
格式如下：
1)僅list(int) dynamic_axes={‘input’:[0,2,3],‘output’:[0,1]}
2)僅dict<int, string> dynamic_axes={‘input’:{0:‘batch’,2:‘height’,3:‘width’},‘output’:{0:‘batch’,1:‘c’}}
3)mixed dynamic_axes={‘input’:{0:‘batch’,2:‘height’,3:‘width’},‘output’:[0,1]}
opset_version——opset的版本，低版本不支援upsample等操作。

import torch
import torch.nn
import onnx

model = torch.load('best.pt')
model.eval()

input_names = ['input']
output_names = ['output']

x = torch.randn(1,3,32,32,requires_grad=True)

torch.onnx.export(model, x, 'best.onnx', input_names=input_names, output_names=output_names, verbose='True')

2 . 執行onnx模型

檢查onnx模型，並使用onnxruntime執行。

import onnx
import onnxruntime as ort

model = onnx.load('best.onnx')
onnx.checker.check_model(model)

session = ort.InferenceSession('best.onnx')
x=np.random.randn(1,3,32,32).astype(np.float32)  # 注意輸入type一定要np.float32!!!!!
# x= torch.randn(batch_size,chancel,h,w)


outputs = session.run(None,input = { 'input' : x })

引數說明：

output_names: default=None
用來指定輸出哪些，以及順序
若為None，則按序輸出所有的output，即返回[output_0,output_1]
若為[‘output_1’,‘output_0’]，則返回[output_1,output_0]
若為[‘output_0’]，則僅返回[output_0:tensor]
input:dict
可以通過session.get_inputs().name獲得名稱
其中key值要求與torch.onnx.export中設定的一致

3.onnx模型輸出與pytorch模型比對

import numpy as np
np.testing.assert_allclose(torch_result[0].detach().numpu(),onnx_result,rtol=0.0001)

如前所述，經驗表明，ONNX 模型的執行效率明顯優於原 PyTorch 模型，這似乎是源於 ONNX 模型生成過程中的優化，這也導致了模型的生成過程比較耗時，但整體效率依舊可觀。

此外，根據對 ONNX 模型和 PyTorch 模型執行結果的統計分析（誤差的均值和標準差），可以看出 ONNX 模型的執行結果誤差很小、基本可靠。

內容參考：https://zhuanlan.zhihu.com/p/422290231

總結

到此這篇關於pytorch模型轉onnx模型的文章就介紹到這了,更多相關pytorch模型轉onnx模型內容請搜尋it145.com以前的文章或繼續瀏覽下面的相關文章希望大家以後多多支援it145.com！