pytorch模型轉onnx模型的方法詳解
學習目標
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模型內容請搜索WalkonNet以前的文章或繼續瀏覽下面的相關文章希望大傢以後多多支持WalkonNet!
推薦閱讀:
- PyTorch模型轉換為ONNX格式實現過程詳解
- 深入學習PyTorch中LSTM的輸入和輸出
- pytorch中常用的乘法運算及相關的運算符(@和*)
- pytorch lstm gru rnn 得到每個state輸出的操作
- pytorch如何獲得模型的計算量和參數量