从零到部署：在Ubuntu 20.04上为YOLOv5模型加速，TensorRT安装与模型转换全流程
从零到部署：在Ubuntu 20.04上为YOLOv5模型加速，TensorRT安装与模型转换全流程
当你的YOLOv5模型在PyTorch环境下运行流畅时，是否思考过如何让它飞得更快？在工业质检、安防监控等实时性要求高的场景中，每毫秒的延迟都可能影响整个系统效能。本文将带你深入TensorRT的世界，从环境配置到模型转换，实现推理速度的质的飞跃。
1. 环境准备与版本匹配
在开始TensorRT之旅前，确保你的Ubuntu 20.04系统已经完成以下基础配置：
- CUDA工具包：推荐11.3版本（与后续TensorRT 8.2完美兼容）
- cuDNN：8.2.x系列（需与CUDA版本严格对应）
- PyTorch：1.12.0+（带GPU支持版本）
验证环境完整性的黄金命令组合：
nvidia-smi # 查看GPU驱动状态 nvcc --version # 确认CUDA编译器版本 python -c "import torch; print(torch.__version__, torch.cuda.is_available())" # 检查PyTorch GPU支持注意：版本不匹配是90%安装失败的根源。建议通过NVIDIA官方文档交叉验证各组件兼容性矩阵。
2. TensorRT的精准安装
2.1 官方包获取与解压
从NVIDIA开发者网站下载对应版本的TensorRT本地安装包（.tar格式），例如：
tar xzvf TensorRT-8.2.5.1.Linux.x86_64-gnu.cuda-11.4.cudnn8.2.tar.gz -C /opt2.2 环境变量配置
在~/.bashrc中添加以下路径（根据实际安装目录调整）：
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/opt/TensorRT-8.2.5.1/lib export PATH=$PATH:/opt/TensorRT-8.2.5.1/bin2.3 Python接口安装
进入解压目录的python文件夹，选择匹配的whl包：
cd /opt/TensorRT-8.2.5.1/python pip install tensorrt-8.2.5.1-cp38-none-linux_x86_64.whl验证安装成功的终极测试：
import tensorrt as trt print(trt.__version__) # 应输出8.2.5.13. YOLOv5模型转换全流程
3.1 权重格式转换（.pt → .wts）
使用tensorrtx项目中的转换脚本：
# 在yolov5项目目录下执行 python gen_wts.py -w yolov5s.pt -o yolov5s.wts关键参数说明：
-w：输入的PyTorch权重文件-o：输出的权重文本文件--img-size：可指定模型输入尺寸（需与训练时一致）
3.2 模型结构适配
修改tensorrtx/yolov5中的配置文件：
- 调整
yololayer.h中的类别数：
static constexpr int CLASS_NUM = 80; // 改为你的实际类别数- 更新
config.h中的输入尺寸：
const static int kInputH = 640; // 高度 const static int kInputW = 640; // 宽度3.3 引擎文件生成
编译转换项目：
mkdir build && cd build cmake .. -DTRT_LIBPATH=/opt/TensorRT-8.2.5.1/lib make -j$(nproc)生成最终引擎文件：
./yolov5 -s yolov5s.wts yolov5s.engine s # 末尾's'对应yolov5s模型4. 性能对比与优化技巧
4.1 基准测试对比
使用相同测试图像进行推理耗时对比：
| 框架类型 | 平均推理时延(ms) | 显存占用(MB) | FPS |
|---|---|---|---|
| PyTorch | 15.2 | 1245 | 65 |
| TensorRT | 4.7 | 873 | 212 |
4.2 高级优化策略
- 精度校准：
from tensorrt import CalibrationAlgoType config.set_flag(trt.BuilderFlag.INT8) config.int8_calibrator = EntropyCalibrator2( data_dir=calib_data_dir, cache_file="yolov5s.calib")- 动态形状支持：
profile = builder.createOptimizationProfile(); profile->setDimensions( "images", OptProfileSelector::kMIN, Dims4(1,3,320,320)); profile->setDimensions( "images", OptProfileSelector::kOPT, Dims4(1,3,640,640));- 层融合优化：
./trtexec --onnx=yolov5s.onnx \ --saveEngine=yolov5s_fp16.engine \ --fp16 \ --workspace=20485. 实战中的避坑指南
- 版本冲突解决方案： 当遇到
undefined symbol错误时，尝试：
patchelf --set-rpath /opt/TensorRT-8.2.5.1/lib libmyplugin.so- 内存不足处理： 在CMakeLists.txt中增加编译选项：
add_definitions(-DUSE_DYNAMIC_CUDA_MEMORY)- 自定义插件开发： 对于特殊算子，需实现IPluginV2接口：
class MyPlugin : public IPluginV2 { const char* getPluginType() const override { return "MyPlugin"; } // 实现其他必要接口... };在模型部署过程中，发现最耗时的环节往往是预处理和后处理而非模型推理本身。通过将图像归一化和NMS操作移植到GPU端执行，整体流水线速度可再提升30%。