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

为为为什么

发布于 2022-08-05 14:35:42

1K0

发布于 2022-08-05 14:35:42

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

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
解压

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

解压后文件夹内文件：

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

安装tensorrt

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

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

安装graphsurgeon wheel

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

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

下载并安装 mmdetection

下载

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

配置环境并安装

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

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

安装 amirstan_plugin

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

更新子模块

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

再次执行

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

新建build文件夹

mkdir build
cd build

生成 makefile

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 文件夹下

编译

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

配置环境变量

export AMIRSTAN_LIBRARY_PATH=<amirstan_plugin_root>/build/lib

安装 mmdetection-to-tensorrt

进入 mmdetection-to-tensorrt 根目录

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：检测有效阈值
配置好后运行该文件即可以生成模型在测试图像上的检测结果