Pytorch模型轉(zhuǎn)onnx模型實(shí)例
如下所示:
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)的是多個(gè)輸入的例子。
在轉(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的錯(cuò)誤。
根據(jù)報(bào)的錯(cuò)誤日志信息打開/home/joy/.tensorflow/venv/lib/python3.6/site-packages/torch/onnx/symbolic_helper.py,在相應(yīng)位置添加print之后,可以定位到具體哪個(gè)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模型實(shí)例就是小編分享給大家的全部內(nèi)容了,希望能給大家一個(gè)參考,也希望大家多多支持腳本之家。
相關(guān)文章
Python實(shí)例解析圖像形態(tài)學(xué)運(yùn)算技術(shù)
形態(tài)學(xué)處理方法是基于對二進(jìn)制圖像進(jìn)行處理的,卷積核決定圖像處理后的效果。本文將為大家詳細(xì)介紹一下OpenCV中的圖像形態(tài)學(xué),感興趣的可以了解一下2022-03-03
Python生命游戲?qū)崿F(xiàn)原理及過程解析(附源代碼)
這篇文章主要介紹了Python生命游戲?qū)崿F(xiàn)原理及過程解析(附源代碼),文中通過示例代碼介紹的非常詳細(xì),對大家的學(xué)習(xí)或者工作具有一定的參考學(xué)習(xí)價(jià)值,需要的朋友可以參考下2019-08-08
python pandas模糊匹配 讀取Excel后 獲取指定指標(biāo)的操作
這篇文章主要介紹了python pandas模糊匹配 讀取Excel后 獲取指定指標(biāo)的操作,具有很好的參考價(jià)值,希望對大家有所幫助。一起跟隨小編過來看看吧2021-03-03
Python?queue雙端隊(duì)列模塊及用法小結(jié)
雙端隊(duì)列是一種具有隊(duì)列和棧性質(zhì)的線性數(shù)據(jù)結(jié)構(gòu),本文主要介紹了Python?queue雙端隊(duì)列模塊及用法小結(jié),文中通過示例代碼介紹的非常詳細(xì),需要的朋友們下面隨著小編來一起學(xué)習(xí)學(xué)習(xí)吧2024-02-02
Python圖像處理庫Pillow的簡單實(shí)現(xiàn)
本文主要介紹了Python圖像處理庫Pillow的簡單實(shí)現(xiàn),文中通過示例代碼介紹的非常詳細(xì),對大家的學(xué)習(xí)或者工作具有一定的參考學(xué)習(xí)價(jià)值,需要的朋友們下面隨著小編來一起學(xué)習(xí)學(xué)習(xí)吧2023-06-06

