将训练好的 mmdetection 模型转为 tensorrt 模型

本文最后更新于：2024年1月14日晚上

mmdetection 是商汤科技（2018 COCO 目标检测挑战赛冠军）和香港中文大学开源的基于Pytorch实现的深度学习目标检测工具箱，性能强大，运算效率高，配置化编程，比较容易训练、测试。但pytorch模型不易于部署，运算速度还有进一步提升的空间，当前比较有效的方法是将模型转换为行为相同的tensorrt模型，本文记录转换流程。

任务思路

转换mmdetection 的 pytorch模型到tensorrt模型有多种方法，本文使用 mmdetection-to-tensorrt 库作为核心，完成直接的模型转换。

该库跳过了通常的 pth -> onnx -> tensorrt 的转换步骤，直接从pth转成tensorrt模型，并且已经成功支持了很多mmdetection 的模型转换。

Support Model/Module

完成步骤

配置环境
安装tensorrt 7.2.3.4
安装 mmdetection-to-tensorrt 库并安装依赖
使用 mmdetection-to-tensorrt 转换模型
结果测试

配置环境

本机 gpu Nvidia GTX 1080 服务器

时间 2021.03

操作系统 Ubuntu 16.04
Nvidia 显卡驱动 460.39
Cuda 版本 11.1
Cudnn 版本 8.1.1

具体配置方法教程很多，在此不再赘述，需要根据个人具体情况配置

安装 tensorrt并配置环境

选择的版本是 tensorrt 7.2.3.4

建议Python环境安装 Anaconda

安装 PyCuda

1	`pip install pycuda`

安装tensorrt

下载tensorrt
- 链接 https://developer.nvidia.com/zh-cn/tensorrt
- 选择 TensorRT-7.2.3.4.Ubuntu-16.04.x86_64-gnu.cuda-11.1.cudnn8.1.tar.gz
解压

1	`tar zxfv TensorRT-7.2.3.4.Ubuntu-16.04.x86_64-gnu.cuda-11.1.cudnn8.1.tar.gz`

解压后文件夹内文件：

1 2	`# ls TensorRT-Release-Notes.pdf bin data doc graphsurgeon include lib onnx_graphsurgeon python samples targets uff`

安装tensorrt

根据自己的 Python 版本选择合适的包进行安装

1 2	`cd TensorRT-7.2.3.4/python pip install tensorrt-7.2.3.4-cp37-none-linux_x86_64.whl`

安装graphsurgeon wheel

1 2	`cd TensorRT-7.2.3.4/graphsurgeon pip install graphsurgeon-0.4.5-py2.py3-none-any.whl`

配置环境变量

export PATH=$PATH:/usr/local/cuda/bin
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda/lib64
export LIBRARY_PATH=$LIBRARY_PATH:/usr/local/cuda/lib64

export PATH=$PATH:"your_path_to_TensorRT-7.2.3.4"
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:"your_path_to_TensorRT-7.2.3.4/lib"

测试

python
import tensorrt
tensorrt.__version__

--> '7.2.3.4'

安装 mmdetection

链接 https://github.com/open-mmlab/mmdetection
安装文档 https://github.com/open-mmlab/mmdetection/blob/master/docs/get_started.md

安装 mmcv

1	`pip install mmcv-full -f https://download.openmmlab.com/mmcv/dist/{cu_version}/{torch_version}/index.html`

下载并安装 mmdetection

下载

1	`git clone git@git.zhlh6.cn:open-mmlab/mmdetection.git`

配置环境并安装

1
2
3

cd mmdetection
pip install -r requirements/build.txt
pip install -v -e .  # or "python setup.py develop"

安装 mmdetection-to-tensorrt

链接 https://github.com/grimoire/mmdetection-to-tensorrt

安装 torch2trt_dynamic

1
2
3

git clone git@git.zhlh6.cn:grimoire/torch2trt_dynamic.git
cd torch2trt_dynamic
python setup.py develop

安装 amirstan_plugin

1 2	`git clone --depth=1 git@git.zhlh6.cn:grimoire/amirstan_plugin.git cd amirstan_plugin`

更新子模块

1	`git submodule update --init --progress --depth=1`

讲道理一句话就可以了，不过我在执行这句命令时报错，如果没报错继续下面的步骤

子模块更新报错解决方案

http协议不好用，需要改成git

修改 amirstan_plugin/.gitmodules 文件

将第三行地址改为 git@github.com:NVIDIA/cub.git

[submodule "third_party/cub"]
	path = third_party/cub
	url = git@github.com:NVIDIA/cub.git 
	branch = 1.8.0

修改 amirstan_plugin/.git/modules/third_party/cub/config

将 remote "origin" 地址改为 git@github.com:NVIDIA/cub.git

[core]
	repositoryformatversion = 0
	filemode = true
	bare = false
	logallrefupdates = true
	worktree = ../../../../third_party/cub
[remote "origin"]
	url = git@github.com:NVIDIA/cub.git 
	fetch = +refs/heads/main:refs/remotes/origin/main
[branch "main"]
	remote = origin
	merge = refs/heads/main

再次执行

1	`git submodule update --init --progress --depth=1`

新建build文件夹

1 2	`mkdir build cd build`

生成 makefile

1	`cmake -DTENSORRT_DIR=${your_path_to_tensorrt} ..`

若输出：

-- Found TensorRT headers at ../TensorRT-7.2.3.4/include
-- Find TensorRT libs at  ../TensorRT-7.2.3.4/lib/libnvinfer.so; ../TensorRT-7.2.3.4/lib/libnvparsers.so; ../TensorRT-7.2.3.4/lib/libnvinfer_plugin.so
-- Found TENSORRT:  ../TensorRT-7.2.3.4/include  
-- WITH_DEEPSTREAM: false
-- GPU_ARCHS is not defined. Generating CUDA code for default SMs: 35;53;61;70;75;80
-- Configuring done
-- Generating done
-- Build files have been written to:  ../amirstan_plugin/build

则说明 makefile 生成成功，保存在 build 文件夹下

编译

1	`make -j10`

此时在build/lib文件夹下生成了很多文件

# ls
libadaptivePoolPlugin_static.a  libcarafeFeatureReassemblePlugin_static.a  libexViewPlugin_static.a             liblayerNormPlugin_static.a     libroiPoolPlugin_static.a         libtorchEmbeddingPlugin_static.a 
libamir_cuda_util.a             libdeformableConvPlugin_static.a           libgridAnchorDynamicPlugin_static.a  libmeshGridPlugin_static.a      libtorchBmmPlugin_static.a        libtorchFlipPlugin_static.a
libamirstan_plugin.so           libdeformablePoolPlugin_static.a           libgridSamplePlugin_static.a         librepeatDimsPlugin_static.a    libtorchCumMaxMinPlugin_static.a  libtorchGatherPlugin_static.a
libbatchedNMSPlugin_static.a    libdelta2bboxPlugin_static.a               libgroupNormPlugin_static.a          libroiExtractorPlugin_static.a  libtorchCumPlugin_static.a        libtorchNMSPlugin_static.a

配置环境变量

1	`export AMIRSTAN_LIBRARY_PATH=<amirstan_plugin_root>/build/lib`

安装 mmdetection-to-tensorrt

进入 mmdetection-to-tensorrt 根目录

1	`python setup.py develop`

测试是否成功

# pip show mmdet2trt

-->
Name: mmdet2trt
Version: 0.3.0
Summary: mmdetection to tensorrt converter
Home-page: UNKNOWN
Author: UNKNOWN
Author-email: UNKNOWN
License: UNKNOWN
Location: /workspace/nfs/tensorrt_test/mmdetection-to-tensorrt
Requires: 
Required-by:

测试

在 mmdetection-to-tensorrt 项目中，运行 demo 文件夹下的 inference.py 文件
修改inference.py 文件中的 parser 参数：
- img：测试图像路径
- config：mmdetection 的模型配置文件
- checkpoint：模型 pth 文件路径
- save_path：tensorrt 模型存放路径
- score-thr：检测有效阈值
配置好后运行该文件即可以生成模型在测试图像上的检测结果