Shortcuts

MMCV中的ONNX Runtime自定义算子

ONNX Runtime介绍

ONNX Runtime是一个跨平台的推理与训练加速器,适配许多常用的机器学习/深度神经网络框架。请访问github了解更多信息。

ONNX介绍

ONNXOpen Neural Network Exchange的缩写,是许多机器学习/深度神经网络框架使用的中间表示(IR)。请访问github了解更多信息。

为什么要在MMCV中添加ONNX自定义算子?

  • 为了验证ONNX模型在ONNX Runtime下的推理的正确性。

  • 为了方便使用了mmcv.ops自定义算子的模型的部署工作。

MMCV已支持的算子

算子 CPU GPU MMCV版本
SoftNMS Y N 1.2.3
RoIAlign Y N 1.2.5
NMS Y N 1.2.7
grid_sampler Y N 1.3.1
CornerPool Y N 1.3.4
cummax Y N 1.3.4
cummin Y N 1.3.4
MMCVModulatedDeformConv2d Y N 1.3.12

如何编译ONNX Runtime自定义算子?

请注意我们仅在onnxruntime>=1.8.1的Linux x86-64 cpu平台上进行过测试

准备工作

  • 克隆代码仓库

git clone https://github.com/open-mmlab/mmcv.git
  • 从ONNX Runtime下载onnxruntime-linuxreleases,解压缩,根据路径创建变量ONNXRUNTIME_DIR并把路径下的lib目录添加到LD_LIBRARY_PATH,步骤如下:

wget https://github.com/microsoft/onnxruntime/releases/download/v1.8.1/onnxruntime-linux-x64-1.8.1.tgz

tar -zxvf onnxruntime-linux-x64-1.8.1.tgz
cd onnxruntime-linux-x64-1.8.1
export ONNXRUNTIME_DIR=$(pwd)
export LD_LIBRARY_PATH=$ONNXRUNTIME_DIR/lib:$LD_LIBRARY_PATH

Linux系统下编译

cd mmcv ## to MMCV root directory
MMCV_WITH_OPS=1 MMCV_WITH_ORT=1 python setup.py develop

如何在python下使用ONNX Runtime对导出的ONNX模型做编译

使用pip安装ONNX Runtime

pip install onnxruntime==1.8.1

推理范例

import os

import numpy as np
import onnxruntime as ort

from mmcv.ops import get_onnxruntime_op_path

ort_custom_op_path = get_onnxruntime_op_path()
assert os.path.exists(ort_custom_op_path)
session_options = ort.SessionOptions()
session_options.register_custom_ops_library(ort_custom_op_path)
## exported ONNX model with custom operators
onnx_file = 'sample.onnx'
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
sess = ort.InferenceSession(onnx_file, session_options)
onnx_results = sess.run(None, {'input' : input_data})

如何为MMCV添加ONNX Runtime的自定义算子

开发前提醒

  • 该算子的ONNX Runtime实现尚未在MMCV中支持已实现算子列表

  • 确保该自定义算子可以被ONNX导出。

添加方法

soft_nms为例:

  1. 在ONNX Runtime头文件目录mmcv/ops/csrc/onnxruntime/下添加头文件soft_nms.h

  2. 在ONNX Runtime源码目录mmcv/ops/csrc/onnxruntime/cpu/下添加算子实现soft_nms.cpp

  3. onnxruntime_register.cpp中注册实现的算子soft_nms

    #include "soft_nms.h"
    
    SoftNmsOp c_SoftNmsOp;
    
    if (auto status = ortApi->CustomOpDomain_Add(domain, &c_SoftNmsOp)) {
    return status;
    }
    
  4. tests/test_ops/test_onnx.py添加单元测试, 可以参考here

最后,欢迎为MMCV添加ONNX Runtime自定义算子 :nerd_face:

已知问题

  • “RuntimeError: tuple appears in op that does not forward tuples, unsupported kind: prim::PythonOp.”

    1. 请注意cummaxcummin算子是在torch >= 1.5.0被添加的。但他们需要在torch version >= 1.7.0才能正确导出。否则会在导出时发生上面的错误。

    2. 解决方法:升级PyTorch到1.7.0以上版本

Read the Docs v: v1.3.18
Versions
latest
stable
v1.3.18
v1.3.17
v1.3.16
v1.3.15
v1.3.14
v1.3.13
Downloads
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.