YOLOV5 -TensorRT linux平台实践
平台配置：
GPU: NVIDIA GeForce 3090
系统：Ubuntu 20.04.4 LTS
必要的包：
tensorrtx容器：hakuyyf/tensorrtx Tags | Docker Hub
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
配置镜像pip源
- cd ~
- mkdir .pip
- vim .pip/pip.conf
在文件中加入清华源：
[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
如果网速太慢，可线下下载.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
其中需要修改：
CUDA_ARCH_BIN ={your GPU Compute Capability}
如果不清楚可执行如下命令：
- python3
- >>> import torch
- >>> print(torch.cuda.get_device_capability())
- >>> (8, 6)
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
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
跑通lenet试验
（1）在pytorchx下生成.wts文件
- cd pytorchx/lenet
# generate lenet.pth
- python3 lenet5.py
# generate lenet5.wts
- python inference.py
一般会出现如下结果：
（2）在tensorrt下生成运行lenet
- cd tensorrtx/lenet
- cp [PATH-OF-pytorchx]/pytorchx/lenet/lenet5.wts .
- mkdir build && cd build
- cmake ..
- make
其中make可能会报错如下：
解决方法：将 void log(Severty severty, const char* msg ) override 更改为：void log(Severty severty, const char* msg ) noexcept
- ./lenet -s
- ./lenet -d
结果如下：
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）下载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
可以在文件夹中看到生成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
最终可见目录如下：
运行默认图片：
./yolov5 -d yolov5s.engine ../samples
（4）使用python下的api进行测试
- cd /home/tensorrtx/yolov5
- python3 yolov5_trt.py
结果如下：
图像实际效果如下：
–END–