Pytorch模型轉(zhuǎn)onnx模型實例
如下所示：
import io
import torch
import torch.onnx
from models.C3AEModel import PlainC3AENetCBAM
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
def test():
model = PlainC3AENetCBAM()
pthfile = r'/home/joy/Projects/models/emotion/PlainC3AENet.pth'
loaded_model = torch.load(pthfile, map_location='cpu')
# try:
# loaded_model.eval()
# except AttributeError as error:
# print(error)
model.load_state_dict(loaded_model['state_dict'])
# model = model.to(device)
#data type nchw
dummy_input1 = torch.randn(1, 3, 64, 64)
# dummy_input2 = torch.randn(1, 3, 64, 64)
# dummy_input3 = torch.randn(1, 3, 64, 64)
input_names = [ "actual_input_1"]
output_names = [ "output1" ]
# torch.onnx.export(model, (dummy_input1, dummy_input2, dummy_input3), "C3AE.onnx", verbose=True, input_names=input_names, output_names=output_names)
torch.onnx.export(model, dummy_input1, "C3AE_emotion.onnx", verbose=True, input_names=input_names, output_names=output_names)
if __name__ == "__main__":
test()
直接將PlainC3AENetCBAM替換成需要轉(zhuǎn)換的模型，然后修改pthfile，輸入和onnx模型名字然后執(zhí)行即可。
注意：上面代碼中注釋的dummy_input2，dummy_input3，torch.onnx.export對應(yīng)的是多個輸入的例子。
在轉(zhuǎn)換過程中遇到的問題匯總
RuntimeError: Failed to export an ONNX attribute, since it's not constant, please try to make things (e.g., kernel size) static if possible
在轉(zhuǎn)換過程中遇到RuntimeError: Failed to export an ONNX attribute, since it's not constant, please try to make things (e.g., kernel size) static if possible的錯誤。
根據(jù)報的錯誤日志信息打開/home/joy/.tensorflow/venv/lib/python3.6/site-packages/torch/onnx/symbolic_helper.py，在相應(yīng)位置添加print之后，可以定位到具體哪個op出問題。
例如：
在相應(yīng)位置添加
print（v.node()）
輸出信息如下：
%124 : Long() = onnx::Gather[axis=0](%122, %121), scope: PlainC3AENetCBAM/Bottleneck[cbam]/CBAM[cbam]/ChannelGate[ChannelGate] # /home/joy/Projects/models/emotion/WhatsTheemotion/models/cbam.py:46:0
原因是pytorch中的tensor.size(1)方式onnx識別不了，需要修改成常量。
以上這篇Pytorch模型轉(zhuǎn)onnx模型實例就是小編分享給大家的全部內(nèi)容了，希望能給大家一個參考，也希望大家多多支持腳本之家。
相關(guān)文章
Python的三元表達式是一種緊湊、簡潔的條件表達式,允許在一行代碼中根據(jù)條件選擇不同的值,三元表達式通常用于需要在單行中根據(jù)條件進行值選擇的情況,有助于提高代碼的可讀性和簡潔性,本文給大家介紹了Python三元表達式的基本用法,需要的朋友可以參考下2023-10-10
本文主要內(nèi)容是python下opencv庫的安裝過程，涉及我在安裝時遇到的問題，并且，將從網(wǎng)上搜集并試用的一些解決方案進行了簡單的匯總，感興趣的小伙伴們可以參考一下2021-06-06
本文主要介紹了如何利用Python制作一個天眼查GUI程序。這次的GUI程序，我們使用的框架是PyQt5，該框架擁有比tkinter更為豐富的內(nèi)置組件，在界面美化方面，貌似也更勝一籌，快跟隨小編一起學起來吧2022-09-09
這篇文章主要介紹了python基礎(chǔ)教程之縮進介紹,本文以if語句為例講解python中的縮進,需要的朋友可以參考下2014-08-08
形態(tài)學處理方法是基于對二進制圖像進行處理的，卷積核決定圖像處理后的效果。本文將為大家詳細介紹一下OpenCV中的圖像形態(tài)學，感興趣的可以了解一下2022-03-03
這篇文章主要介紹了 PyCharm中如何創(chuàng)建帶有注釋的py文件,本文給大家介紹的非常詳細,對大家的學習或工作具有一定的參考借鑒價值,需要的朋友參考下吧2024-08-08
Python生命游戲?qū)崿F(xiàn)原理及過程解析（附源代碼）
這篇文章主要介紹了Python生命游戲?qū)崿F(xiàn)原理及過程解析（附源代碼）,文中通過示例代碼介紹的非常詳細，對大家的學習或者工作具有一定的參考學習價值,需要的朋友可以參考下2019-08-08
python pandas模糊匹配 讀取Excel后 獲取指定指標的操作
這篇文章主要介紹了python pandas模糊匹配 讀取Excel后 獲取指定指標的操作，具有很好的參考價值，希望對大家有所幫助。一起跟隨小編過來看看吧2021-03-03
雙端隊列是一種具有隊列和棧性質(zhì)的線性數(shù)據(jù)結(jié)構(gòu),本文主要介紹了Python?queue雙端隊列模塊及用法小結(jié),文中通過示例代碼介紹的非常詳細,需要的朋友們下面隨著小編來一起學習學習吧2024-02-02
本文主要介紹了Python圖像處理庫Pillow的簡單實現(xiàn)，文中通過示例代碼介紹的非常詳細，對大家的學習或者工作具有一定的參考學習價值，需要的朋友們下面隨著小編來一起學習學習吧2023-06-06