Skip to content

Latest commit

 

History

History
129 lines (86 loc) · 4.96 KB

onnxruntime_op.md

File metadata and controls

129 lines (86 loc) · 4.96 KB
Raw

MMCV中的ONNX Runtime自定义算子

ONNX Runtime介绍

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

ONNX介绍

ONNX是Open 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-linux：releases，解压缩，根据路径创建变量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自定义算子 🤓

已知问题

  • "RuntimeError: tuple appears in op that does not forward tuples, unsupported kind: prim::PythonOp."
    1. 请注意cummax和cummin算子是在torch >= 1.5.0被添加的。但他们需要在torch version >= 1.7.0才能正确导出。否则会在导出时发生上面的错误。
    2. 解决方法：升级PyTorch到1.7.0以上版本

引用