1 onnxruntime
Onnx runtime是一个跨平台的机器学习模型加速器,可以在不同的硬件和操作系统上运行,可以加载和推理任意机器学习框架导出的onnx模型并进行加速。
如要使用onnxruntime,一般通过以下步骤:
- 从机器学习框架中将模型导出为onnx
- 使用onnxruntime加载onnx模型并进行推理
onnxruntime官网:https://onnxruntime.ai/
Github地址:https://github.com/microsoft/onnxruntime
1.1 onnxruntime安装
onnxruntime在python上有两个版本:cpu和gpu版本,在一个python环境中只能安装一个版本,gpu版本包含了大部分cpu版本的内容,所以在有gpu的情况下,尽量安装gpu版本。
cpu版本
pip install onnxruntime
gpu版本
pip install onnxruntime-gpu
1.2 从pytorch导出onnx模型
使用pytorch
的torch.onnx.export
函数导出onnx模型,这里以pytorch的resnet18预训练模型为例
# -*- coding: utf-8 -*-
import torch
import torch.nn as nn
import torchvision.models
import onnx
import onnxruntime
if __name__ == '__main__':
resnet18 = torchvision.models.resnet18(pretrained=True)
input_image = torch.rand([1, 3, 224, 224],dtype=torch.float)
onnx_outpath = 'resnet18.onnx'
torch.onnx.export(resnet18,
input_image,
onnx_outpath,
opset_version=13,
verbose=True,
do_constant_folding=True,
input_names=['input'],
output_names=['output']
)
# 检查导出的onnx模型
onnx_model = onnx.load(onnx_outpath)
onnx.checker.check_model(onnx_model, full_check=True)
inferred = onnx.shape_inference.infer_shapes(onnx_model, check_type=True)
1.3 使用onnxruntime对onnx模型进行推理
从pytorch导出onnx模型之后,就可以使用onnxruntime加载模型并进行推理,还是以resnet18为例,示例代码如下
# -*- coding: utf-8 -*-
import torch
import torch.nn as nn
import torchvision.models
import onnx
import onnxruntime
if __name__ == '__main__':
resnet18 = torchvision.models.resnet18(pretrained=True)
input_image = torch.rand([1, 3, 224, 224],dtype=torch.float)
onnx_outpath = 'resnet18.onnx'
torch.onnx.export(resnet18,
input_image,
onnx_outpath,
opset_version=13,
verbose=True,
do_constant_folding=True,
input_names=['input'],
output_names=['output']
)
# 检查导出的onnx模型
onnx_model = onnx.load(onnx_outpath)
onnx.checker.check_model(onnx_model, full_check=True)
inferred = onnx.shape_inference.infer_shapes(onnx_model, check_type=True)
# 使用onnxruntime对onnx模型进行推理
providers = ["CUDAExecutionProvider"]
ort_session = onnxruntime.InferenceSession(onnx_outpath, providers=providers)
output = ort_session.run(None,{'input': input_image.numpy()})
result = output[0]
print(result)
1.4 onnxruntime API
onnxruntime的python版本的API文档地址为:https://onnxruntime.ai/docs/api/python/api_summary.html,有兴趣的可以仔细看看,下面简要介绍onnxruntime中经常使用的api。
1.4.1 onnxruntime.InferenceSession
1.4.1.1 onnxruntime.InferenceSession类
类原型
class onnxruntime.InferenceSession(path_or_bytes, sess_options=None, providers=None, provider_options=None, **kwargs)
类初始化参数
- path_or_bytes:onnx或者ort模型的文件名或者序列化模型
- sess_options:会话选项
- providers:按优先级递减顺序排列的可选程序序列,如
CUDAExecutionProvider
,CPUExecutionProvider
,具体的可参考Execution Providers - provider_options:与providers列出的程序相对应的可选字典序列
1.4.1.2 onnxruntime.InferenceSession.run成员函数
函数原型
run(output_names, input_feed, run_options=None)
函数参数
- output_names:输出的名称
- input_feed:格式为
{input_name:input_value}
的字典 - run_options:参考onnxruntime.RunOptions
函数返回值
返回一个列表,列表中的每一个输出结果为numpy数组、sparse tensor(稀疏张量)、列表或者字典。
本文作者:StubbornHuang
版权声明:本文为站长原创文章,如果转载请注明原文链接!
原文标题:Python – 使用onnxruntime加载和推理onnx模型
原文链接:https://www.stubbornhuang.com/2427/
发布于:2022年11月30日 13:59:20
修改于:2023年06月21日 17:48:48
声明:本站所有文章,如无特殊说明或标注,均为本站原创发布。任何个人或组织,在未征得本站同意时,禁止复制、盗用、采集、发布本站内容到任何网站、书籍等各类媒体平台。如若本站内容侵犯了原著者的合法权益,可联系我们进行处理。
评论
51