一、將訓(xùn)練好的模型轉(zhuǎn)換格式為ONNX格式

例如pytorch模型轉(zhuǎn)換：

def torch2onnx(model, save_path):

"""

:param model:

:param save_path: XXX/XXX.onnx

:return:

"""

model.eval()

data = torch.rand(1, 3, 224, 224)

input_names = ["input"]

output_names = ["out"]

torch.onnx._export(model, data, save_path, export_params=True, opset_version=11, input_names=input_names, output_names=output_names)

print("torch2onnx finish.")

支持動(dòng)態(tài)形狀的輸入和輸出：

def torch2onnx_dynamic(model, save_path):

"""

:param model:

:param save_path: XXX/XXX.onnx

:return:

"""

model.eval()

data = torch.rand(1, 3, 224, 224)

input_names = ["input"] # ncnn需要

output_names = ["out"] # ncnn需要

torch.onnx._export(model, data, save_path, export_params=True, opset_version=11, input_names=input_names,

output_names=output_names, dynamic_axes={'input': [2, 3], 'out': [2, 3]})

print("torch2onnx finish.")

二、安裝onnxruntime

注意：onnxruntime-gpu版本在0.4以上時(shí)需要CUDA 10

pip install onnxruntime

pip install onnxruntime-gpu

onnxruntime幫助文檔：

https://microsoft.github.io/onnxruntime/python/tutorial.html

三、onnxruntime使用方法

加載模型：

session = onnxruntime.InferenceSession("./dmnet.onnx")

加載圖片：

img = cv2.imread(img_path)

img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

tensor = transforms.ToTensor()(img)

tensor = tensor.unsqueeze_(0)

執(zhí)行推理：

注意：這里的"input"是和轉(zhuǎn)onnx格式時(shí)的名字對(duì)應(yīng)的。

result = session.run([], {"input": tensor.cpu().numpy()})

創(chuàng)作挑戰(zhàn)賽新人創(chuàng)作獎(jiǎng)勵(lì)來(lái)咯，堅(jiān)持創(chuàng)作打卡瓜分現(xiàn)金大獎(jiǎng)

總結(jié)

以上是生活随笔為你收集整理的onnx模型部署 python_onnxruntime模型部署流程的全部?jī)?nèi)容，希望文章能夠幫你解決所遇到的問(wèn)題。

如果覺(jué)得生活随笔網(wǎng)站內(nèi)容還不錯(cuò)，歡迎將生活随笔推薦給好友。