Python – 使用onnxruntime加载和推理onnx模型

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(稀疏张量)、列表或者字典。