liu-mengyang / trt-fairmot

该项目实现了对以DLA34作为骨干网络的FairMOT的TensorRT迁移加速,同时这也是一个NVIDIA-阿里云的Hackathon2021竞赛项目。
16 stars 4 forks source link

trt-fairmot

该项目实现了多目标跟踪算法FairMOT的TensorRT版本,是DLA34骨干网络版本。

最新动态

可运行的环境

本项目推荐在Linux系统下使用tensorrt:21.02-py3版本的NGC TensorRT Docker容器运行,此外你需要一块算力较强的GPU显卡作为硬件设备。

本README所展示加速效果的硬件测试平台

一些重要的软件环境

选择的技术路线

基于ONNX的方案

仍然存在的问题

由于本项目还不是特别成熟,仍然存在一些问题,所以把这一小节放到了相对靠前的位置。

能够达到的效果

针对预测网络

针对预测网络,本项目能够在fp32精度下以与原版模型几乎为0的相对误差实现x1.48的单视频输入加速比。在fp16精度下针对三个分支分别会有大约2%、2%和54%的相对误差,这里的误差较大,具体原因还要进一步排查,但它能带来x2.36的单视频输入加速比。

针对整个系统

目前,由于缺少对于预测网络预处理与后处理关于TensorRT的优化,系统直接使用TensorRT模型进行预测会因为数据的前后传输与类型转换带来额外开销导致整体性能下降。

安装

  1. 部署运行环境 请跟随NGC安装教程部署能够运行本项目的容器镜像并运行容器环境。该容器环境会直接包含你所需要的Python和TensorRT环境,此外按照该教程安装一些必要的库。

  2. 克隆项目 在磁盘的任一位置克隆本项目。

    git clone https://github.com/liu-mengyang/trt-fairmot
    cd trt-fairmot
  3. 安装本项目使用的一些第三方库

    pip3 install -r requirements.txt
    apt install libgl1-mesa-glx
  4. 下载原版权重 这里给出原版仓库所给出的一些链接 fairmot_dla34.pth [Google] [Baidu, 提取码:uouv] [Onedrive]。下载完成后将权重文件放到weights文件夹下。

    mv fairmot_dla34.pth ./weights/

使用

进入TensorRT_ONNX_impl目录。

cd TensorRT_ONNX_impl

构建TensorRT Engine

  1. 编译Plugin

    make -j
  2. 导出PyTorch模型为ONNX模型

    python build_onnx_engine.py
    python replace_dcn_plugin.py
  3. 将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

MOT样例

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"

MOT Benchmark测试

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"

其他

待办清单

本项目会一直得到维护,目前的不足会在未来被加以完善。

相关的项目

PyTorch版本

Plugin开发