1. 环境准备与基础概念在开始构建YOLOv5的TensorRT引擎之前我们需要先准备好开发环境。我推荐使用Python 3.7或3.8版本因为TensorRT对这些版本的兼容性最好。硬件方面你需要一块支持CUDA的NVIDIA显卡我测试用的是GTX 1650对于YOLOv5s这样的轻量级模型已经足够。首先安装必要的依赖包pip install torch1.7.0 torchvision0.8.1 pip install onnx1.10.0 pip install tensorrt8.2.4.2这里有个小技巧安装TensorRT时建议直接从NVIDIA官网下载对应版本的whl文件进行安装这样可以避免很多兼容性问题。我遇到过直接用pip安装最新版导致API不匹配的情况折腾了好久才发现是版本问题。TensorRT的核心价值在于它能够对神经网络进行深度优化。它会自动完成层融合Layer Fusion、精度校准Precision Calibration、内核自动调优Kernel Auto-Tuning等优化操作。想象一下这就像是一个经验丰富的厨师不仅知道如何做菜还能根据食材特性选择最合适的烹饪方式让最终出品又快又好。2. 导出ONNX模型YOLOv5官方仓库已经提供了非常方便的导出脚本我们只需要运行export.py即可。但在此之前我建议先检查一下模型结构import torch model torch.load(yolov5s.pt)[model] print(model.model[-1]) # 查看最后的Detect层导出ONNX的命令很简单python export.py --weights yolov5s.pt --include onnx --img 640 --batch 1这里有几个关键参数需要注意--img 640指定输入图像尺寸--batch 1设置批处理大小为1--dynamic如果需要动态输入尺寸可以加上这个参数我在第一次导出时遇到了一个坑输出的ONNX模型在Netron中查看时发现输出维度是[1,25200,85]而不是预期的三个检测头。这是因为YOLOv5在导出模式下会自动将三个头的输出合并。如果你需要保留原始输出结构需要修改models/yolo.py中的Detect类。3. TensorRT核心组件解析构建TensorRT引擎主要涉及以下几个核心对象我把它比喻成汽车制造流程Builder- 就像汽车工厂的厂长负责统筹整个建造过程Network- 这是汽车的蓝图定义了模型的结构Config- 相当于生产工艺配置决定如何优化这个模型Parser- 就像翻译官把ONNX格式的模型转换成TensorRT能理解的格式先来看如何创建Builder和Networkimport tensorrt as trt logger trt.Logger(trt.Logger.WARNING) builder trt.Builder(logger) network builder.create_network(1 int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))这里有个重要细节EXPLICIT_BATCH标志。TensorRT处理批次数据有两种模式显式批次更适合现代CNN模型。我曾在某个项目中去掉了这个标志结果模型推理完全错误debug了整整一天才发现是这个原因。4. 解析ONNX模型有了Network之后我们就可以用ONNX解析器来填充网络定义了parser trt.OnnxParser(network, logger) with open(yolov5s.onnx, rb) as model: if not parser.parse(model.read()): for error in range(parser.num_errors): print(parser.get_error(error))解析过程中最常见的错误是某些算子不被支持。YOLOv5的Focus层在早期TensorRT版本中就有这个问题。解决方法要么是修改模型结构要么升级TensorRT版本。我在实际项目中就遇到过这种情况最终选择将Focus层替换为等价的卷积操作。解析成功后我们可以检查网络的输入输出for i in range(network.num_inputs): input network.get_input(i) print(fInput {i}: {input.name}, shape{input.shape}, dtype{input.dtype}) for i in range(network.num_outputs): output network.get_output(i) print(fOutput {i}: {output.name}, shape{output.shape}, dtype{output.dtype})5. 构建优化配置Config对象决定了TensorRT如何优化我们的模型。对于YOLOv5最重要的几个配置是config builder.create_builder_config() config.max_workspace_size 1 30 # 1GB的临时空间 config.set_flag(trt.BuilderFlag.FP16) # 开启FP16模式关于workspace size我的经验是设置太小会导致某些优化无法进行太大又浪费内存。对于YOLOv5s1GB通常足够。如果是更大的YOLOv5x模型可能需要增加到2GB。FP16模式可以显著提升推理速度但对精度的影响需要评估。我在实际测试中发现YOLOv5使用FP16后推理速度提升约2倍而mAP仅下降不到1%这个trade-off非常值得。6. 构建并序列化引擎一切准备就绪后就可以构建引擎了engine builder.build_engine(network, config) with open(yolov5s.engine, wb) as f: f.write(engine.serialize())这个过程可能会花费几分钟时间具体取决于模型复杂度和硬件性能。构建过程中TensorRT会尝试各种优化策略包括层融合比如ConvBNReLU合并为一个操作内核选择为每个操作选择最优的CUDA内核内存优化减少数据传输和内存占用我曾经对比过同一模型在TensorRT和原生PyTorch上的性能在GTX 1650上TensorRT版本的推理速度提升了3-4倍内存占用减少了约40%。这种优化效果在边缘设备上尤为明显。7. 验证引擎性能构建完成后我们可以加载引擎进行测试with open(yolov5s.engine, rb) as f, trt.Runtime(logger) as runtime: engine runtime.deserialize_cuda_engine(f.read()) context engine.create_execution_context() # 分配输入输出内存 inputs, outputs, bindings [], [], [] for binding in engine: size trt.volume(engine.get_binding_shape(binding)) dtype trt.nptype(engine.get_binding_dtype(binding)) # 分配设备内存 device_mem cuda.mem_alloc(size * dtype.itemsize) bindings.append(int(device_mem)) if engine.binding_is_input(binding): inputs.append({mem: device_mem, shape: engine.get_binding_shape(binding)}) else: outputs.append({mem: device_mem, shape: engine.get_binding_shape(binding)})测试推理速度时建议使用timeit进行多次测量因为第一次推理通常会有额外的初始化开销。我在Jetson Nano上测试时发现连续推理100次的平均时间比单次推理要稳定和准确得多。8. 常见问题与解决方案在实际项目中我遇到过不少坑这里分享几个典型问题的解决方法不支持的算子如果遇到类似Unsupported operator: GridSample的错误可以考虑以下方案升级TensorRT到最新版本修改模型结构用支持的算子替代自定义插件实现该算子精度问题FP16模式下可能出现检测框偏移这时可以检查模型训练时是否使用了混合精度尝试在Config中禁用FP16对关键层强制使用FP32精度性能不如预期如果优化后的速度提升不明显可以检查是否真的使用了TensorRT引擎有时会意外fallback到原生模式尝试不同的工作空间大小使用trtexec工具进行基准测试对比记得有一次客户报告模型推理结果完全错误最后发现是因为他们在导出ONNX时使用了动态尺寸但在构建TensorRT引擎时却指定了固定尺寸。这种维度不匹配的问题往往会导致难以察觉的错误。