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

本文最后更新于:2022年8月10日 上午

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

任务思路

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

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

Support Model/Module
  • Faster R-CNN
  • Cascade R-CNN
  • Double-Head R-CNN
  • Group Normalization
  • Weight Standardization
  • DCN
  • SSD
  • RetinaNet
  • Libra R-CNN
  • FCOS
  • Fovea
  • CARAFE
  • FreeAnchor
  • RepPoints
  • NAS-FPN
  • ATSS
  • PAFPN
  • FSAF
  • GCNet
  • Guided Anchoring
  • Generalized Attention
  • Dynamic R-CNN
  • Hybrid Task Cascade
  • DetectoRS
  • Side-Aware Boundary Localization
  • YOLOv3
  • PAA
  • CornerNet(WIP)
  • Generalized Focal Loss
  • Grid RCNN
  • VFNet
  • GROIE
  • Mask R-CNN(experiment)
  • Cascade Mask R-CNN(experiment)
  • Cascade RPN

完成步骤

  • 配置环境
  • 安装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

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
  • 配置环境变量
1
2
3
4
5
6
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"
  • 测试
1
2
3
4
5
python
import tensorrt
tensorrt.__version__

--> '7.2.3.4'

安装 mmdetection

安装 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

安装 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

1
2
3
4
[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

1
2
3
4
5
6
7
8
9
10
11
12
[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} ..

若输出:

1
2
3
4
5
6
7
8
-- 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文件夹下生成了很多文件

1
2
3
4
5
# 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
  • 测试是否成功
1
2
3
4
5
6
7
8
9
10
11
12
13
# 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:检测有效阈值
  • 配置好后运行该文件即可以生成模型在测试图像上的检测结果

参考资料


将训练好的 mmdetection 模型转为 tensorrt 模型
https://www.zywvvd.com/notes/study/deep-learning/deploy/tensorrt-model-gen-mmdection/tensorrt-model-gen/
作者
Yiwei Zhang
发布于
2021年3月17日
许可协议