YOLOV5 -TensorRT linux平台实践_tensorrt exec

YOLOV5 -TensorRT linux平台实践

平台配置：

GPU: NVIDIA GeForce 3090

系统：Ubuntu 20.04.4 LTS

必要的包：

tensorrtx:GitHub - wang-xinyu/tensorrtx: Implementation of popular deep learning networks with TensorRT network definition API

tensorrtx容器：hakuyyf/tensorrtx Tags | Docker Hub

pytorchx:GitHub - wang-xinyu/pytorchx: Implement popular deep learning networks in pytorch, used by tensorrtx.

yolov5:https://github.com/ultralytics/yolov5.git

预训练模型：yolov5s.pt

1、环境配置

拉取容器并进入容器

- docker pull hakuyyf/tensorrtx:trt8_cuda11
- docker run --gpus all --name [容器名称] -itd -p [物理机端口]:[容器内部端口] -v [物理机地址]:[镜像地址] --privileged=true --ipc=host [IMAGE ID]  /bin/bash
- docker exec -it [CONTAINER ID] /bin/bash
- apt update
1
2
3
4

配置镜像pip源

- cd ~
- mkdir .pip 
- vim .pip/pip.conf
1
2
3

在文件中加入清华源：
[global]
index-url=https://pypi.tuna.tsinghua.edu.cn/simple
[install]
use-mirrors=true
mirrors=https://pypi.tuna.tsinghua.edu.cn/simple
trusted-host=pypi.tuna.tsinghua.edu.cn

可参考该网址进行配置：(46条消息) pip源_淘小欣的博客-CSDN博客_pip源

安装torch

需去官网查询不同cuda版本对应的torch版本

copy 上述安装命令进行安装

- pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu116
1

如果网速太慢，可线下下载.whl文件进行安装

备注：如果直接pip安装，可能出现如下错误：

下载并编译opencv

- cd /home 
- git clone https://github.com/opencv/opencv.git
- git clone https://github.com/opencv/opencv_contrib.git 
- cd opencv
- git checkout 4.5.5
- cd ../
- cd opencv_contrib
- git checkout 4.5.5
- cd ../
- cd opencv
- mkdir build && cd build
- cmake -D CMAKE_BUILD_TYPE=RELEASE -D CMAKE_INSTALL_PREFIX=/usr/local -D INSTALL_PYTHON_EXAMPLES=ON -D INSTALL_C_EXAMPLES=OFF -D OPENCV_ENABLE_NONFREE=ON -D WITH_CUDA=ON -D WITH_CUDNN=ON -D OPENCV_DNN_CUDA=ON -D ENABLE_FAST_MATH=1 -D CUDA_FAST_MATH=1 -D CUDA_ARCH_BIN=8.6 -D WITH_CUBLAS=1 -D OPENCV_EXTRA_MODULES_PATH=/home/opencv_contrib/modules/ -D HAVE_opencv_python3=ON -D PYTHON_EXECUTABLE=/usr/bin/python3 -D BUILD_EXAMPLES=ON ..
- make -j
- make install
- ldconfig
- ln -s /usr/local/lib/python3.8/site-packages/cv2/python-3.8/cv2.cpython-38-x86_64-linux-gnu.so /usr/lib/python3.8/cv2.so
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

其中需要修改：

CUDA_ARCH_BIN ={your GPU Compute Capability}

如果不清楚可执行如下命令：

- python3
- >>> import torch
- >>> print(torch.cuda.get_device_capability())
- >>> (8, 6)

1
2
3
4
5

OPENCV_EXTRA_MODULES_PATH=/home/opencv_contrib/modules/

配置成功的标志是：

下载yolov5-5.0并配置其环境

- cd /home
- git clone -b v5.0 https://github.com/ultralytics/yolov5.git
- cd yolov5
- pip install -r requirements.txt
1
2
3
4

2、在tensorrtx下测试lenet

下载tensorrtx和pytorchx

- cd /home
- git clone -b yolov5-v4.0 https://github.com/wang-xinyu/tensorrtx.git
- git clone https://github.com/wang-xinyu/pytorchx.git
1
2
3

跑通lenet试验

（1）在pytorchx下生成.wts文件

- cd pytorchx/lenet
# generate lenet.pth
- python3 lenet5.py
# generate lenet5.wts
- python inference.py
1
2
3
4
5

一般会出现如下结果：

（2）在tensorrt下生成运行lenet

- cd tensorrtx/lenet
- cp [PATH-OF-pytorchx]/pytorchx/lenet/lenet5.wts .
- mkdir build && cd build
- cmake ..
- make
1
2
3
4
5

其中make可能会报错如下：

解决方法：将 void log(Severty severty, const char* msg ) override 更改为：void log(Severty severty, const char* msg ) noexcept

- ./lenet -s
- ./lenet -d
1
2

结果如下：

3、TensorRT下yolov5-5.0运行

本次运行都是基于默认配置，如需使用自己数据进行配置，请参考：/tensorrtx/yolov5/README

- Choose the model n/s/m/l/x/n6/s6/m6/l6/x6 from command line arguments.
- Input shape defined in yololayer.h
- Number of classes defined in yololayer.h, **DO NOT FORGET TO ADAPT THIS, If using your own model**
- INT8/FP16/FP32 can be selected by the macro in yolov5.cpp, **INT8 need more steps, pls follow `How to Run` first and then go the `INT8 Quantization` below**
- GPU id can be selected by the macro in yolov5.cpp
- NMS thresh in yolov5.cpp
- BBox confidence thresh in yolov5.cpp
- Batch size in yolov5.cpp
1
2
3
4
5
6
7
8

（1）下载yolov5-5.0对应的预训练模型yolov5s.pth

地址如下：https://github.com/ultralytics/yolov5/releases/download/v5.0/yolov5s.pt

将该预训练模型放在如下位置：/home/pre_train/yolov5s.pt

（2）生成yolov5模型对应的.wts文件

- cd /home/yolov5
- cp ../tensorrtx/yolov5/gen_wts.py .
- python gen_wts.py -w ../pre_train/yolov5s.pt -o yolov5s.wts
1
2
3

可以在文件夹中看到生成yolov5s.wts

（3）在tensortrx下生成engine并运行

生成yolov5s.engine

- cd /home/tensorrtx/yolov5
- mkdir build && cd build
- cp ../../../yolov5/yolov5s.wts .
- cmake ..
- make 
- ./yolov5 -s yolov5s.wts yolov5s.engine
1
2
3
4
5
6

最终可见目录如下：

运行默认图片：

 ./yolov5 -d yolov5s.engine ../samples
1

（4）使用python下的api进行测试

- cd /home/tensorrtx/yolov5
- python3  yolov5_trt.py
1
2

结果如下：

图像实际效果如下：

–END–