该项目实现了多目标跟踪算法FairMOT的TensorRT版本,是DLA34骨干网络版本。
本项目推荐在Linux系统下使用tensorrt:21.02-py3版本的NGC TensorRT Docker容器运行,此外你需要一块算力较强的GPU显卡作为硬件设备。
本README所展示加速效果的硬件测试平台
一些重要的软件环境
基于ONNX的方案
trtexec
工具将已开发的自定义算子与导出的ONNX模型集成生成一个TensorRT Engine。
TRTAPI_impl
的工程文件夹,在其中我们选择了一条动态构建的基于TensorRT Python API
的技术路线。这条基于API构建TRT Engine的路线需要更多的时间,但同时也会取得更好的加速效果。然而很遗憾,目前本项目关于这一路线的实现还有一些问题存在,它得到的结果与原模型的结果有较大差异,还不能使用。不过没关系,这个项目会一直得到维护,并在未来完成这一技术路线的实现。由于本项目还不是特别成熟,仍然存在一些问题,所以把这一小节放到了相对靠前的位置。
针对预测网络
针对预测网络,本项目能够在fp32精度下以与原版模型几乎为0的相对误差实现x1.48的单视频输入加速比。在fp16精度下针对三个分支分别会有大约2%、2%和54%的相对误差,这里的误差较大,具体原因还要进一步排查,但它能带来x2.36的单视频输入加速比。
针对整个系统
目前,由于缺少对于预测网络预处理与后处理关于TensorRT的优化,系统直接使用TensorRT模型进行预测会因为数据的前后传输与类型转换带来额外开销导致整体性能下降。
部署运行环境 请跟随NGC安装教程部署能够运行本项目的容器镜像并运行容器环境。该容器环境会直接包含你所需要的Python和TensorRT环境,此外按照该教程安装一些必要的库。
克隆项目 在磁盘的任一位置克隆本项目。
git clone https://github.com/liu-mengyang/trt-fairmot
cd trt-fairmot
安装本项目使用的一些第三方库
pip3 install -r requirements.txt
apt install libgl1-mesa-glx
下载原版权重
这里给出原版仓库所给出的一些链接 fairmot_dla34.pth
[Google] [Baidu, 提取码:uouv] [Onedrive]。下载完成后将权重文件放到weights
文件夹下。
mv fairmot_dla34.pth ./weights/
进入TensorRT_ONNX_impl目录。
cd TensorRT_ONNX_impl
编译Plugin
make -j
导出PyTorch模型为ONNX模型
python build_onnx_engine.py
python replace_dcn_plugin.py
将ONNX模型转化为TensorRT Engine
sh build_trt.sh
sh build_trt_fp16.sh
mv fairmot.trt ../weights
mv fairmot_fp16.trt ../weights
compare_onnx_fairmot.py
文件对PyTorch版本和TensorRT版本的FairMOT进行了运行并比较了性能,它会输出fp32和fp16的TRT Engine所能带来的加速比和相对误差。
python compare_onnx_fairmot.py
Demo.py
是一个执行多目标跟踪任务的样例程序,通过执行它可以看到针对一个测试视频进行处理后得到的结果。同时也可以比较出PyTorch版本和TensorRT版本下的整体性能差异。
# 默认执行PyTorch版本
python Demo.py
# 执行TensorRT版本,默认执行fp32精度版本
python Demo.py --trt_enable==True
# 执行fp16精度版本的TensorRT Engine
python Demo.py --trt_enable==True --trt_load=="../weights/fairmot_fp16.trt"
test_benchmark.py
负责benchmark的测试工作,这里本项目选择的MOT17数据集,在MOTChallenge下载这一数据集,本项目默认数据集保存在/workspace/dataset/
下。
# 默认进行原版的benchmark测试
python test_benchmark.py
# 进行TensorRT fp32版本的benchmark测试
python test_benchmark.py --trt_enable==True
# 进行TensorRT fp16版本的benchmark测试
python test_benchmark.py --trt_enable==True --trt_load=="../weights/fairmot_fp16.trt"
Export
函数即可,并且仅需对应好层和输入输出对象,而不必纠结于权重文件的命名。然而很遗憾,目前API搭建的工作存在一些还未排查出来的问题,但未来会解决它。本项目会一直得到维护,目前的不足会在未来被加以完善。
PyTorch版本
Plugin开发