快捷導(dǎo)航

pytorch模型部署到onnx的詳細(xì)過程

更新時間：2023年08月28日 10:25:53 作者：Liang-ml

這篇文章主要介紹了如何簡單的將pytorch模型部署到onnx,本文結(jié)合示例代碼給大家介紹的非常詳細(xì),對大家的學(xué)習(xí)或工作具有一定的參考借鑒價值,需要的朋友可以參考下

1. 創(chuàng)建一個pytorch模型

這里我用的U2Net，直接加載好訓(xùn)練出的權(quán)重

model = U2Net(class_nums=4)
model.load_state_dict(torch.load(checkpoint_path))

2. 將pytorch模型轉(zhuǎn)成onnx格式

x = torcg.randn(1,3,512,512)
with torch.no_grad():
torch.onnx.export(
model,
x,
"onnx_file_name.onnx",
opset_version=11,
input_names=['input'],
output_names=['output']
)

torch.onnx.export()函數(shù)用于將pytorch模型轉(zhuǎn)成onnx格式的函數(shù)，其中的參數(shù)如下

model：需要轉(zhuǎn)換的模型，加載好權(quán)重
args：模型的任意一組輸入，注意維度
f：導(dǎo)出的onnx模型的文件名
opset_version：表示ONNX算子集的版本，隨著深度學(xué)習(xí)的發(fā)展，新算子會不斷誕生，算子集就需要擴(kuò)充，截至20230824，官網(wǎng)已經(jīng)發(fā)布20個版本
input_names：輸入tensor的名稱
output_names：輸出tensor的名稱

轉(zhuǎn)換成功會在指定的路徑f下生成onnx模型文件

可用onnx中的函數(shù)驗證模型文件是否正確

import onnx
onnx_model = onnx.load("onnx_file_name.onnx")
try:
onnx.checker.check_model(onnx_model)
except Exception as ex:
print(f'ERROR: {ex}')
else:
print(f'model correct')

可以使用netron查看轉(zhuǎn)換的onnx模型的詳細(xì)內(nèi)容鏈接 https://netron.app/

點(diǎn)擊某個算子節(jié)點(diǎn)，可以查看算子的具體信息，每個算子記錄了算子的屬性、圖結(jié)構(gòu)和權(quán)重三類信息

算子屬性信息(attributes)，對于卷積來說，算子屬性包含了卷積核的大小、步長等屬性
圖結(jié)構(gòu)信息(node properties)，算子節(jié)點(diǎn)在計算圖中的名稱、鄰邊的信息，如上圖節(jié)點(diǎn)名為conv2，輸入數(shù)據(jù)叫683，權(quán)重叫1324
權(quán)重信息指算子存儲的網(wǎng)絡(luò)權(quán)重信息，點(diǎn)擊w后的+號可以查看權(quán)重信息的具體內(nèi)容

3. 推理引擎ONNXRuntime

使用onnxruntime可以運(yùn)行輸出的onnx模型，代碼如下

input_img 是一個與onnx模型輸入維度一致的歸一化的np.array
model = onnxruntime.InferenceSession(onnx_file)
inputs = {'input':input_img}
output = model.run(['output'],inputs)[0]

onnxruntime.InferenceSession用于獲取一個ONNX Runtime的推理器，其參數(shù)為用于推理的ONNX模型文件

run方法用于模型推理，第一個參數(shù)為輸出張量名的列表，第二個參數(shù)為輸入值的字典，key為導(dǎo)出模型時設(shè)置的張量名，value為輸入張量

到此這篇關(guān)于簡單的將pytorch模型部署到onnx的文章就介紹到這了,更多相關(guān)pytorch模型部署到onnx內(nèi)容請搜索腳本之家以前的文章或繼續(xù)瀏覽下面的相關(guān)文章希望大家以后多多支持腳本之家！

您可能感興趣的文章: