先要安装好显卡驱动、CUDA、CUDNN
以ubuntu20.04、显卡1650安装470版本的显卡驱动、11.3版本的CUDA及8.2版本的CUDNN为例
进入网站:
https://developer.nvidia.com/nvidia-tensorrt-8x-download
进行勾选下载:
TAR是免安装直接解压可用的
解压:
tar -zxvf TensorRT-8.4.2.4.Linux.x86_64-gnu.cuda-11.6.cudnn8.4.tar.gz
cd TensorRT-8.4.2.4/samples/sampleMNIST
make
cd ../../bin
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/YourPath/TensorRT-8.4.2.4/lib
./sample_mnist
终端打印出如下内容表明cuda+cudnn+tensorrt安装正常:
可以在.bashrc里面加入TensorRT的路径:
# TensorRT
export TRT_PATH=/usr/local/TensorRT-8.4.2.4
export PATH=$PATH:$TRT_PATH/bin
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$TRT_PATH/lib
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$TRT_PATH/targets/x86_64-linux-gnu/lib
先下载
https://github.com/triple-Mu/YOLOv8-TensorRT.git
git clone https://github.com/triple-Mu/YOLOv8-TensorRT.git
cd YOLOv8-TensorRT
pip install -r requirements.txt
C++ build:
CMakeLists中需要将TensorRT路径改一下
# TensorRT
set(TensorRT_INCLUDE_DIRS /usr/include/x86_64-linux-gnu)
set(TensorRT_LIBRARIES /usr/lib/x86_64-linux-gnu)
加下来就可以编译了
export root=${PWD}
cd csrc/detect/normal
mkdir build
cmake ..
make
mv yolov8 ${root}
cd ${root}
官方给出的model是pt格式,我们需要TensorRT要用的engine格式,
PyTorch model -> ONNX -> TensorRT Engine
Then,transform!
# PyTorch model -> ONNX
from ultralytics import YOLO
# Load a model
model = YOLO("yolov8s.pt") # load a pretrained model (recommended for training)
success = model.export(format="onnx", opset=11, simplify=True) # export the model to onnx format
assert success
ONNX -> TensorRT Engine需要用TensorRT
终端进入TensorRT目录
cd samples/trtexec
make
然后退到TensorRT进入bin文件夹可以看到trtexec可执行文件
/bin/trtexec \
--onnx=yolov8s.onnx \
--saveEngine=yolov8s.engine \
--fp16
路径不要错就可以成功将onnx转为engine格式了,运行时间可能会有些长
如果在C++ build时遇到下面error可以按下面的解决方案:
error: invalid initialization of reference of type ‘std::vector<cv::String>&’ from expression of type ‘std::vector<std::__cxx11::basic_string<char> >’
69 | cv::glob(path + "/*.jpg", imagePathList);
将第46行:
std::vector<std::string> imagePathList;
改为:
std::vector<cv::string> imagePathList;
接下来就可以使用了
# infer image
./yolov8 yolov8s.engine data/bus.jpg
# infer images
./yolov8 yolov8s.engine data
# infer video
./yolov8 yolov8s.engine data/test.mp4 # the video path
如:
./yolov8 pt/yolov8n.engine data/zidane.jpg