1. 創(chuàng)建一個(gè)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ā)展，新算子會(huì)不斷誕生，算子集就需要擴(kuò)充，截至20230824，官網(wǎng)已經(jīng)發(fā)布20個(gè)版本
• input_names：輸入tensor的名稱
• output_names：輸出tensor的名稱

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

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

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

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

3. 推理引擎ONNXRuntime

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

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

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

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

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

最新評(píng)論