从零到部署在Ubuntu 20.04上为YOLOv5模型加速TensorRT安装与模型转换全流程当你的YOLOv5模型在PyTorch环境下运行流畅时是否思考过如何让它飞得更快在工业质检、安防监控等实时性要求高的场景中每毫秒的延迟都可能影响整个系统效能。本文将带你深入TensorRT的世界从环境配置到模型转换实现推理速度的质的飞跃。1. 环境准备与版本匹配在开始TensorRT之旅前确保你的Ubuntu 20.04系统已经完成以下基础配置CUDA工具包推荐11.3版本与后续TensorRT 8.2完美兼容cuDNN8.2.x系列需与CUDA版本严格对应PyTorch1.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)FPSPyTorch15.2124565TensorRT4.78732124.2 高级优化策略精度校准from tensorrt import CalibrationAlgoType config.set_flag(trt.BuilderFlag.INT8) config.int8_calibrator EntropyCalibrator2( data_dircalib_data_dir, cache_fileyolov5s.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 --onnxyolov5s.onnx \ --saveEngineyolov5s_fp16.engine \ --fp16 \ --workspace20485. 实战中的避坑指南版本冲突解决方案 当遇到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%。