一、部署环境

  • TensorRT 8.6

  • CUDA 11.7

  • TensorRTx

二、TensorRTx 核心优势

  • 针对自定义网络做了比原生 TensorRT 更深度的优化。

  • 在 NVIDIA 4090 上,YOLOv5s 的 TensorRTx 推理速度比原生 TensorRT 快 15-25%,FP16 精度下可达 350+ FPS(可参考该性能表现,YOLOv8 部署后性能优势一致)。

三、跨平台 / 跨版本兼容性

TensorRTx 封装了版本适配逻辑,同一套代码可兼容 TensorRT 8.0-8.6,且提供了跨平台的 CUDA 核函数实现,无需额外修改代码即可适配不同环境。

四、部署步骤

4.1 克隆 TensorRTx 仓库

git clone https://github.com/wang-xinyu/tensorrtx.git

4.2 下载 YOLOv8 权重文件

# 进入 YOLOv8 部署目录
cd tensorrtx/yolov8
# 下载 YOLOv8n.pt(轻量版,适合快速测试;也可替换为 s/m/l/x 版本)
wget https://github.com/ultralytics/assets/releases/download/v8.2.0/yolov8n.pt

4.3 将 .pt 权重转为 .wts 格式

python gen_wts.py -w yolov8n.pt -o yolov8n.wts -t detect

4.4 配置参数并将 .wts 权重转为 .engine 格式

4.4.1 修改配置文件

打开路径~/yolov8_tensorrt/yolov8/include/config.h,根据自身需求修改以下参数(默认适配 COCO 数据集,无需额外修改可直接使用):

  • 默认 kNumClass=80(对应 COCO 数据集 80 类别),若使用自定义数据集,需修改为自身数据集的类别数。

  • 输入尺寸默认 kInputH=640、kInputW=640,匹配 YOLOv8 默认输入尺寸,无需修改。

  • 默认量化方式为 FP16,适配 4090 性能最优,无需修改。

4.4.2 修改 CMakeLists.txt 并编译
# 修改 CMakeLists.txt 中 TensorRT 路径(根据自身 TensorRT 安装路径调整)
# 编译构建
mkdir build
cd build
cp yolov8.wts build
cmake ..
make
4.4.3 链接库并执行转换
# 链接 TensorRT 库(根据自身 TensorRT 安装路径调整)
export LD_LIBRARY_PATH=/usr/local/TensorRT-8.4.3.1/lib/:$LD_LIBRARY_PATH
# 执行权重转换(生成 yolov8.engine 文件)
./yolov8_det -s yolov8n.wts yolov8.engine n
Logo

腾讯云面向开发者汇聚海量精品云计算使用和开发经验,营造开放的云计算技术生态圈。

更多推荐