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模型

使用pytorchtorch.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:按优先级递减顺序排列的可选程序序列,如CUDAExecutionProviderCPUExecutionProvider,具体的可参考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(稀疏张量)、列表或者字典。