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版本
<code class="language-python line-numbers">pip install onnxruntime </code>
gpu版本
<code class="language-python line-numbers">pip install onnxruntime-gpu </code>
1.2 从pytorch导出onnx模型
使用pytorch
的torch.onnx.export
函数导出onnx模型,这里以pytorch的resnet18预训练模型为例
<code class="language-python line-numbers"># -*- 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) </code>
1.3 使用onnxruntime对onnx模型进行推理
从pytorch导出onnx模型之后,就可以使用onnxruntime加载模型并进行推理,还是以resnet18为例,示例代码如下
<code class="language-python line-numbers"># -*- 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) </code>
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类
类原型
<code class="language-python line-numbers">class onnxruntime.InferenceSession(path_or_bytes, sess_options=None, providers=None, provider_options=None, **kwargs) </code>
类初始化参数
- path_or_bytes:onnx或者ort模型的文件名或者序列化模型
- sess_options:会话选项
- providers:按优先级递减顺序排列的可选程序序列,如
CUDAExecutionProvider
,CPUExecutionProvider
,具体的可参考Execution Providers - provider_options:与providers列出的程序相对应的可选字典序列
1.4.1.2 onnxruntime.InferenceSession.run成员函数
函数原型
<code class="language-python line-numbers">run(output_names, input_feed, run_options=None) </code>
函数参数
- 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
声明:本站所有文章,如无特殊说明或标注,均为本站原创发布。任何个人或组织,在未征得本站同意时,禁止复制、盗用、采集、发布本站内容到任何网站、书籍等各类媒体平台。如若本站内容侵犯了原著者的合法权益,可联系我们进行处理。
评论
49