MMCV中的ONNX Runtime自定义算子¶
为什么要在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的自定义算子¶
添加方法¶
以soft_nms
为例:
在ONNX Runtime头文件目录
mmcv/ops/csrc/onnxruntime/
下添加头文件soft_nms.h
在ONNX Runtime源码目录
mmcv/ops/csrc/onnxruntime/cpu/
下添加算子实现soft_nms.cpp
在onnxruntime_register.cpp中注册实现的算子
soft_nms
#include "soft_nms.h" SoftNmsOp c_SoftNmsOp; if (auto status = ortApi->CustomOpDomain_Add(domain, &c_SoftNmsOp)) { return status; }
在
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
.”请注意
cummax
和cummin
算子是在torch >= 1.5.0被添加的。但他们需要在torch version >= 1.7.0才能正确导出。否则会在导出时发生上面的错误。解决方法:升级PyTorch到1.7.0以上版本