Skip to content

Based on TensorRT v8.2, build network for YOLOv5-v5.0 by myself, speed up YOLOv5-v5.0 inferencing

License

Notifications You must be signed in to change notification settings

emptysoal/TensorRT-v8-YOLOv5-v5.0

Repository files navigation

TensorRT v8.2 加速部署 YOLOv5-v5.0

项目简介

  • 使用 TensorRT 原生API构建 YOLO 网络,将 PyTorch 模型转为.plan 序列化文件,加速模型推理;
  • 基于 TensorRT 8.2.4 版本,具体环境见下方的环境构建部分;
  • 主要参考 tensorrtx 项目,但作者本人根据自己编程习惯,做了大量改动;
  • 未使用Cuda加速图像预处理的项目链接:no_cuda_preproc

项目特点

  • 这里对比和 tensorrtx 项目中 YOLOv5-v5.0 的不同,并不是说孰优孰劣,只是有些地方更符合作者个人习惯
tensorrtx 本项目 备注
1 implicit(隐式 batch) explicit(显式 batch) 此不同为最大的不同,代码中很多的差异都源于此
2 Detect Plugin 继承自 IPluginV2IOExt Detect Plugin 继承自 IPluginV2DynamicExt
3 Detect Plugin 被编译为动态链接库 Detect Plugin 直接编译到最终的可执行文件
4 异步推理(context.enqueue) 同步推理(context.executeV2) 作者亲测在速度方面无差别,同步写法更简便
5 INT8量化时,采用OpenCV的dnn模块将图像转换为张量 INT8量化时,自定义的方法将图像转换为张量
6 C++加opencv实现预处理 cuda编程实现预处理加速 v5.0之后的版本也有,两种不同的实现

除上述外,还有很多其他编码上的不同,不一一赘述。

推理速度

  • 基于GPU:GeForce RTX 2080 Ti
FP32 FP16 INT8
6 ms 3 ms 3 ms

备注:本项目的推理时间包括:预处理、前向传播、后处理,tensorrtx 项目仅计算了前向传播时间。

环境构建

宿主机基础环境

  • Ubuntu 16.04
  • GPU:GeForce RTX 2080 Ti
  • docker,nvidia-docker

基础镜像拉取

docker pull nvcr.io/nvidia/tensorrt:22.04-py3
  • 该镜像中各种环境版本如下:
CUDA cuDNN TensorRT python
11.6.2 8.4.0.27 8.2.4.2 3.8.10

安装其他库

  1. 创建 docker 容器

    docker run -it --gpus device=0 --shm-size 32G -v /home:/workspace nvcr.io/nvidia/tensorrt:22.04-py3 bash

    其中-v /home:/workspace将宿主机的/home目录挂载到容器中,方便一些文件的交互,也可以选择其他目录

    • 将容器的源换成国内源
    cd /etc/apt
    rm sources.list
    vim sources.list
    • 将下面内容拷贝到文件sources.list
    deb http://mirrors.aliyun.com/ubuntu/ focal main restricted universe multiverse
    deb http://mirrors.aliyun.com/ubuntu/ focal-security main restricted universe multiverse
    deb http://mirrors.aliyun.com/ubuntu/ focal-updates main restricted universe multiverse
    deb http://mirrors.aliyun.com/ubuntu/ focal-proposed main restricted universe multiverse
    deb http://mirrors.aliyun.com/ubuntu/ focal-backports main restricted universe multiverse
    deb-src http://mirrors.aliyun.com/ubuntu/ focal main restricted universe multiverse
    deb-src http://mirrors.aliyun.com/ubuntu/ focal-security main restricted universe multiverse
    deb-src http://mirrors.aliyun.com/ubuntu/ focal-updates main restricted universe multiverse
    deb-src http://mirrors.aliyun.com/ubuntu/ focal-proposed main restricted universe multiverse
    deb-src http://mirrors.aliyun.com/ubuntu/ focal-backports main restricted universe multiverse
    • 更新源
    apt update
  2. 安装 OpenCV-4.5.0

    • OpenCV-4.5.0源码链接如下,下载 zip 包,解压后放到宿主机/home目录下,即容器的/workspace目录下
    https://github.com/opencv/opencv
    • 下面操作均在容器中
    # 安装依赖
    apt install build-essential
    apt install libgtk2.0-dev pkg-config libavcodec-dev libavformat-dev libswscale-dev
    apt install libtbb2 libtbb-dev libjpeg-dev libpng-dev libtiff-dev libdc1394-22-dev
    # 开始安装 OpenCV
    cd /workspace/opencv-4.5.0
    mkdir build
    cd build
    cmake -D CMAKE_INSTALL_PREFIX=/usr/local -D CMAKE_BUILD_TYPE=Release -D OPENCV_GENERATE_PKGCONFIG=ON -D OPENCV_ENABLE_NONFREE=True ..
    make -j6
    make install

运行项目

  1. 获取 .wts 文件
  • 主要过程为:把本项目的pth2wts.py文件复制到官方yolov5-v5.0目录下,在官方yolov5-v5.0目录下运行 python pth2wts.py,得到para.wts文件
  • 具体过程可参考下面步骤
git clone -b v5.0 https://github.com/ultralytics/yolov5.git
git clone https://github.com/emptysoal/yolov5-v5.0_tensorrt-v8.2.git
# download https://github.com/ultralytics/yolov5/releases/download/v5.0/yolov5s.pt
cp {tensorrt}/pth2wts.py {ultralytics}/yolov5
cd {ultralytics}/yolov5
python pth2wts.py
# a file 'para.wts' will be generated.
  1. 构建 .plan 序列化文件并推理
  • 主要过程为:把上一步生成的para.wts文件复制到本项目目录下,在本项目中依次运行make./trt_infer
  • 具体过程可参考下面步骤
cp {ultralytics}/yolov5/para.wts {tensorrt}/
cd {tensorrt}/
mkdir images  # and put some images in it
# update CLASS_NUM in yololayer.h if your model is trained on custom dataset
# you can also update INPUT_H、INPUT_W in yololayer.h, update NET(s/m/l/x) in trt_infer.cpp
make
./trt_infer
# result images will be generated in present dir