MXNET模型进行部署时可以使用TensorRT,其过程通常为将MXNET模型转换为ONNX,再通过TensorRT的ONNX解释器转换为TensorRT的序列化文件。但此过程中,MXNET的OP不一定可以转换为ONNX的结果,而ONNX的OP也只有部分可以在TensorRT中实现。
TensorRT支持操作
可参照TensorRT官方文档,这里面除了介绍支持的OP操作外,还有关于支持的平台和支持int8压缩的层等信息。
可以看到,由于ONNX开源,应该可以找到通过重新编译添加支持的操作,但未找到添加plugin的说明,从issues上看有人提这个问题,但目前没有得到回复。
ONNX支持操作
ONNX支持的操作可以在github的ONNX项目中查看,Operators.md中写明了支持的操作及其说明。
按照TensorRT 4 基本介绍及TensorFlow模型转化Uff格式导入(二)的说法,ONNX是不支持自定义层的。
MXNET支持操作
MXNET支持的操作可以在MXNET官网的NDArray API中查看,也写命令支持的操作及其说明。
MXNET实际可以转换为ONNX的操作可以在github 的python/mxnet/contrib/onnx/mx2onnx/_op_translations.py中查看,需要注意的是这里面出现的操作只能证明是通过master分支编译的版本可以支持,而自己电脑中安装的版本需要在python路径相应位置下查看。比如master分支的作为1.5.0未发布版,新增了对Deconvolution等操作的支持,而当前发布的到1.4.1不支持这些操作,如果只是临时转换模型,可以将master分支下的contrib/onnx整体替换自己python版本的文件。
查找问题解答
MXNET转TensorRT遇到问题是很正常的事,可以多去github的issue和相关论坛查找下:
1、MXNET的issue
2、MXNET的论坛(有时候可能打不开…)
3、ONNX的issue
4、TensorRT的论坛