097、TensorRT 部署 YOLOONNX到TRT 引擎构建到Context 推理到序列化和反序列化上周帮同事排查一个YOLOv8部署到Jetson Orin上的问题模型跑起来检测框全飘了坐标完全不对。他用的就是网上抄的一段TensorRT推理代码看着挺完整但就是结果不对。我让他把engine序列化出来反序列化后打印了一下网络输入输出的name发现他构建engine时用的输入名是“images”但推理时写死了“input”TensorRT根本找不到这个tensor直接用了未初始化的内存。这种坑我踩过不止一次今天就把整个流程从ONNX到TRT引擎构建、Context推理、序列化反序列化每一步的细节和容易翻车的地方都写清楚。从ONNX开始别以为导出就完事了YOLO模型训练完第一步是导出ONNX。很多人直接torch.onnx.export一把梭结果部署时各种报错。这里有个关键点动态batch和静态batch的选择。如果你部署的场景batch size固定为1强烈建议导出静态batch的ONNXTensorRT对静态shape的优化更激进推理速度能快10%-15%。动态batch虽然灵活但TensorRT在构建引擎时会为每个可能的batch size都做优化构建时间翻倍而且模型体积更大。导出时别忘了设置opset_version11以上TensorRT对低版本opset支持不太好。我习惯用opset_version17兼容性最好。还有一个容易忽略的参数dynamic_axes。如果你确实需要动态batch这样写dynamic_axes{images:{0:batch_size},output:{0:batch_size}}注意这里的key要和模型forward的输入输出参数名一致别自己瞎起名字。我见过有人把输入叫“input.1”输出叫“169”这种名字在TensorRT里调试起来能让人崩溃。构建TRT引擎Builder、Config、Network三件套拿到ONNX后构建引擎是核心步骤。TensorRT的Builder就像个工厂Config是工厂的配置单Network是你要生产的产品图纸。这三者的关系必须理清。先创建builderimporttensorrtastrt TRT_LOGGERtrt.Logger(trt.Logger.WARNING)buildertrt.Builder(TRT_LOGGER)这里日志级别我一般用WARNINGINFO会打印太多构建细节生产环境没必要。但调试阶段可以改成INFO能看清每一步优化做了什么。创建network时有个参数很多人搞错networkbuilder.create_network(1int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))这个EXPLICIT_BATCH标志必须加否则ONNX解析器会报错。不加的话TensorRT默认使用隐式batch模式而ONNX导出时用的是显式batch两者不兼容。这个坑我踩了整整一天才找到原因。然后创建config设置工作空间大小configbuilder.create_builder_config()config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE,130)# 1GB工作空间大小不是越大越好但太小会导致TensorRT无法做某些优化。对于YOLOv8s这种模型1GB足够了。如果是YOLOv8x建议给到2GB。注意这个值受GPU显存限制别设得比显存还大builder会直接崩掉。接下来解析ONNXparsertrt.OnnxParser(network,TRT_LOGGER)withopen(onnx_path,rb)asf:ifnotparser.parse(f.read()):forerrorinrange(parser.num_errors):print(parser.get_error(error))raiseRuntimeError(ONNX解析失败)解析失败时一定要打印所有错误信息TensorRT的错误提示虽然有时候很抽象但至少能告诉你哪个op出了问题。常见错误是某些算子不支持比如GridSample在旧版TensorRT上就不支持需要自己写plugin。构建引擎耐心等待或者加点加速构建引擎是最耗时的步骤对于YOLOv8s在RTX 3090上大概需要30秒到1分钟。如果是在Jetson上可能得等5分钟以上。serialized_enginebuilder.build_serialized_network(network,config)这里我直接用了build_serialized_network它返回的是序列化后的engine字节流可以直接保存到文件。如果你用build_engine返回的是engine对象还需要手动序列化。一步到位更省事。构建过程中TensorRT会做层融合、精度校准、内存优化等操作。如果你看到日志里出现[W]开头的警告比如某个层被回退到FP32别慌这通常不影响最终结果只是性能可能不如预期。反序列化从文件到可执行引擎保存engine到文件withopen(yolov8s.engine,wb)asf:f.write(serialized_engine)反序列化时注意要创建新的runtimeruntimetrt.Runtime(TRT_LOGGER)withopen(yolov8s.engine,rb)asf:engineruntime.deserialize_cuda_engine(f.read())这里有个细节反序列化后的engine只能在相同硬件和相同TensorRT版本上使用。你把RTX 3090上构建的engine拿到Jetson Orin上直接报错。所以生产环境通常是在目标设备上构建engine或者用Docker统一环境。创建Context真正的推理执行者Engine是静态的模型描述Context才是真正执行推理的上下文。每个Context可以独立执行推理互不干扰。contextengine.create_execution_context()如果你需要多batch推理记得设置动态shape的实际大小context.set_binding_shape(0,(batch_size,3,640,640))这个步骤很容易被忽略。如果你导出ONNX时用了动态batch但推理时没设置binding shapeTensorRT会默认使用你在ONNX里指定的最小shape导致输入数据尺寸不匹配结果全错。推理流程内存管理是重灾区推理的核心是数据搬运从CPU到GPU执行推理再从GPU到CPU。这里最容易出问题的是内存分配和释放。先分配GPU内存importpycuda.driverascudaimportpycuda.autoinit# 获取输入输出binding的索引input_idxengine.get_binding_index(images)output_idxengine.get_binding_index(output)# 获取binding的shapeinput_shapeengine.get_binding_shape(input_idx)output_shapeengine.get_binding_shape(output_idx)# 分配GPU内存d_inputcuda.mem_alloc(input_shape.numel()*np.dtype(np.float32).itemsize)d_outputcuda.mem_alloc(output_shape.numel()*np.dtype(np.float32).itemsize)# 创建CUDA流streamcuda.Stream()这里注意get_binding_index的参数名必须和ONNX里的输入输出名完全一致。我习惯在构建engine后先打印所有binding的名字foriinrange(engine.num_bindings):print(fBinding{i}:{engine.get_binding_name(i)})这样能避免名字写错。执行推理# 准备输入数据input_datapreprocess(image)# 假设已经预处理成(1,3,640,640)的numpy数组# 拷贝输入到GPUcuda.memcpy_htod_async(d_input,input_data,stream)# 执行推理context.execute_async_v2(bindings[int(d_input),int(d_output)],stream_handlestream.handle)# 拷贝输出回CPUoutput_datanp.empty(output_shape,dtypenp.float32)cuda.memcpy_dtoh_async(output_data,d_output,stream)# 同步流stream.synchronize()这里有个关键点execute_async_v2的bindings参数是一个列表顺序必须和engine中binding的顺序一致。如果你不确定顺序可以用engine.get_binding_name(i)逐个确认。我见过有人把输入和输出顺序搞反结果模型输出了一堆NaN。后处理YOLO特有的解码TensorRT的输出是原始的张量对于YOLOv8输出shape通常是(1, 84, 8400)其中84是4个坐标80个类别概率8400是预测框数量。你需要自己实现解码defpostprocess(output,conf_threshold0.5,iou_threshold0.45):# output shape: (1, 84, 8400)outputoutput.squeeze(0)# (84, 8400)outputoutput.transpose(1,0)# (8400, 84)boxes[]scores[]class_ids[]foriinrange(output.shape[0]):class_scoresoutput[i,4:]max_scoreclass_scores.max()ifmax_scoreconf_threshold:continueclass_idclass_scores.argmax()# 解码坐标这里假设是cx,cy,w,h格式cx,cy,w,houtput[i,:4]x1cx-w/2y1cy-h/2x2cxw/2y2cyh/2boxes.append([x1,y1,x2,y2])scores.append(max_score)class_ids.append(class_id)# NMSindicescv2.dnn.NMSBoxes(boxes,scores,conf_threshold,iou_threshold)# 处理结果...注意这里的坐标是相对于输入图像尺寸的需要缩放到原图尺寸。很多人忘记这一步导致检测框位置完全不对。序列化和反序列化的最佳实践在实际部署中engine的构建和推理通常是分开的。构建一次engine保存为文件后续推理时直接加载。但这里有个坑engine文件依赖具体的GPU型号和驱动版本。如果你在开发机上构建engine然后拷贝到生产服务器很可能加载失败。我的做法是在部署脚本中先检查是否存在engine文件如果不存在则构建存在则直接加载。同时在engine文件名中加入模型名称和输入尺寸的哈希值避免混淆。importhashlibdefget_engine_path(model_name,input_shape):hash_strf{model_name}_{input_shape}hash_valhashlib.md5(hash_str.encode()).hexdigest()[:8]returnf{model_name}_{hash_val}.engine另外序列化后的engine文件是二进制格式不要用文本编辑器打开也不要在不同字节序的机器间传输。我见过有人把engine文件用git管理结果每次pull下来都损坏就是因为git自动转换了换行符。性能调优别让GPU闲着如果你发现推理速度不达标可以从几个方向排查使用FP16或INT8精度在config中设置config.set_flag(trt.BuilderFlag.FP16)对于YOLO模型FP16几乎不影响精度但速度能提升30%-50%。INT8需要校准数据集麻烦一些但速度提升更明显。减少CPU-GPU数据传输如果输入数据是连续帧可以复用GPU内存避免每次推理都重新分配。我通常会在初始化时分配好内存推理时只做memcpy。使用多流并行如果你的应用需要同时处理多个视频流可以为每个流创建独立的Context和CUDA流实现并行推理。注意Context不是线程安全的每个线程必须有自己的Context。检查是否被CPU瓶颈拖累有时候预处理和后处理比推理本身还慢。考虑用GPU做预处理比如用CUDA核函数做resize和归一化。个人经验总结做了这么多年YOLO部署踩过的坑能写一本书。最后给几个实在的建议永远不要相信ONNX导出后的模型。导出后先用ONNX Runtime跑一遍确认结果正确再交给TensorRT。很多问题出在ONNX导出这一步TensorRT只是背锅侠。构建engine时加--verbose日志。虽然慢一点但能看清每一步优化做了什么出了问题好定位。保存engine文件时记录构建参数。我习惯在文件名里包含模型名、输入尺寸、精度模式、TensorRT版本号。这样即使过了半年看到文件名就知道这个engine是怎么来的。别在生产环境用动态batch。除非你真的需要否则静态batch省心省力。动态batch的构建时间、内存占用、推理延迟都不可控。测试时用真实数据。不要用全零或随机数据测试推理流程YOLO对输入数据很敏感全零输入可能输出正常但真实图片就崩了。用一张真实的测试图片从预处理到后处理全流程跑通。最后也是最重要的做好异常处理。TensorRT的API调用失败时有时候不会抛异常而是默默返回空指针。每个API调用后都要检查返回值别偷懒。这套流程我用了三年从YOLOv5到YOLOv8从Jetson Nano到A100基本没出过大问题。希望你能少走弯路。
097、TensorRT 部署 YOLO：ONNX到TRT 引擎构建到Context 推理到序列化和反序列化
097、TensorRT 部署 YOLOONNX到TRT 引擎构建到Context 推理到序列化和反序列化上周帮同事排查一个YOLOv8部署到Jetson Orin上的问题模型跑起来检测框全飘了坐标完全不对。他用的就是网上抄的一段TensorRT推理代码看着挺完整但就是结果不对。我让他把engine序列化出来反序列化后打印了一下网络输入输出的name发现他构建engine时用的输入名是“images”但推理时写死了“input”TensorRT根本找不到这个tensor直接用了未初始化的内存。这种坑我踩过不止一次今天就把整个流程从ONNX到TRT引擎构建、Context推理、序列化反序列化每一步的细节和容易翻车的地方都写清楚。从ONNX开始别以为导出就完事了YOLO模型训练完第一步是导出ONNX。很多人直接torch.onnx.export一把梭结果部署时各种报错。这里有个关键点动态batch和静态batch的选择。如果你部署的场景batch size固定为1强烈建议导出静态batch的ONNXTensorRT对静态shape的优化更激进推理速度能快10%-15%。动态batch虽然灵活但TensorRT在构建引擎时会为每个可能的batch size都做优化构建时间翻倍而且模型体积更大。导出时别忘了设置opset_version11以上TensorRT对低版本opset支持不太好。我习惯用opset_version17兼容性最好。还有一个容易忽略的参数dynamic_axes。如果你确实需要动态batch这样写dynamic_axes{images:{0:batch_size},output:{0:batch_size}}注意这里的key要和模型forward的输入输出参数名一致别自己瞎起名字。我见过有人把输入叫“input.1”输出叫“169”这种名字在TensorRT里调试起来能让人崩溃。构建TRT引擎Builder、Config、Network三件套拿到ONNX后构建引擎是核心步骤。TensorRT的Builder就像个工厂Config是工厂的配置单Network是你要生产的产品图纸。这三者的关系必须理清。先创建builderimporttensorrtastrt TRT_LOGGERtrt.Logger(trt.Logger.WARNING)buildertrt.Builder(TRT_LOGGER)这里日志级别我一般用WARNINGINFO会打印太多构建细节生产环境没必要。但调试阶段可以改成INFO能看清每一步优化做了什么。创建network时有个参数很多人搞错networkbuilder.create_network(1int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))这个EXPLICIT_BATCH标志必须加否则ONNX解析器会报错。不加的话TensorRT默认使用隐式batch模式而ONNX导出时用的是显式batch两者不兼容。这个坑我踩了整整一天才找到原因。然后创建config设置工作空间大小configbuilder.create_builder_config()config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE,130)# 1GB工作空间大小不是越大越好但太小会导致TensorRT无法做某些优化。对于YOLOv8s这种模型1GB足够了。如果是YOLOv8x建议给到2GB。注意这个值受GPU显存限制别设得比显存还大builder会直接崩掉。接下来解析ONNXparsertrt.OnnxParser(network,TRT_LOGGER)withopen(onnx_path,rb)asf:ifnotparser.parse(f.read()):forerrorinrange(parser.num_errors):print(parser.get_error(error))raiseRuntimeError(ONNX解析失败)解析失败时一定要打印所有错误信息TensorRT的错误提示虽然有时候很抽象但至少能告诉你哪个op出了问题。常见错误是某些算子不支持比如GridSample在旧版TensorRT上就不支持需要自己写plugin。构建引擎耐心等待或者加点加速构建引擎是最耗时的步骤对于YOLOv8s在RTX 3090上大概需要30秒到1分钟。如果是在Jetson上可能得等5分钟以上。serialized_enginebuilder.build_serialized_network(network,config)这里我直接用了build_serialized_network它返回的是序列化后的engine字节流可以直接保存到文件。如果你用build_engine返回的是engine对象还需要手动序列化。一步到位更省事。构建过程中TensorRT会做层融合、精度校准、内存优化等操作。如果你看到日志里出现[W]开头的警告比如某个层被回退到FP32别慌这通常不影响最终结果只是性能可能不如预期。反序列化从文件到可执行引擎保存engine到文件withopen(yolov8s.engine,wb)asf:f.write(serialized_engine)反序列化时注意要创建新的runtimeruntimetrt.Runtime(TRT_LOGGER)withopen(yolov8s.engine,rb)asf:engineruntime.deserialize_cuda_engine(f.read())这里有个细节反序列化后的engine只能在相同硬件和相同TensorRT版本上使用。你把RTX 3090上构建的engine拿到Jetson Orin上直接报错。所以生产环境通常是在目标设备上构建engine或者用Docker统一环境。创建Context真正的推理执行者Engine是静态的模型描述Context才是真正执行推理的上下文。每个Context可以独立执行推理互不干扰。contextengine.create_execution_context()如果你需要多batch推理记得设置动态shape的实际大小context.set_binding_shape(0,(batch_size,3,640,640))这个步骤很容易被忽略。如果你导出ONNX时用了动态batch但推理时没设置binding shapeTensorRT会默认使用你在ONNX里指定的最小shape导致输入数据尺寸不匹配结果全错。推理流程内存管理是重灾区推理的核心是数据搬运从CPU到GPU执行推理再从GPU到CPU。这里最容易出问题的是内存分配和释放。先分配GPU内存importpycuda.driverascudaimportpycuda.autoinit# 获取输入输出binding的索引input_idxengine.get_binding_index(images)output_idxengine.get_binding_index(output)# 获取binding的shapeinput_shapeengine.get_binding_shape(input_idx)output_shapeengine.get_binding_shape(output_idx)# 分配GPU内存d_inputcuda.mem_alloc(input_shape.numel()*np.dtype(np.float32).itemsize)d_outputcuda.mem_alloc(output_shape.numel()*np.dtype(np.float32).itemsize)# 创建CUDA流streamcuda.Stream()这里注意get_binding_index的参数名必须和ONNX里的输入输出名完全一致。我习惯在构建engine后先打印所有binding的名字foriinrange(engine.num_bindings):print(fBinding{i}:{engine.get_binding_name(i)})这样能避免名字写错。执行推理# 准备输入数据input_datapreprocess(image)# 假设已经预处理成(1,3,640,640)的numpy数组# 拷贝输入到GPUcuda.memcpy_htod_async(d_input,input_data,stream)# 执行推理context.execute_async_v2(bindings[int(d_input),int(d_output)],stream_handlestream.handle)# 拷贝输出回CPUoutput_datanp.empty(output_shape,dtypenp.float32)cuda.memcpy_dtoh_async(output_data,d_output,stream)# 同步流stream.synchronize()这里有个关键点execute_async_v2的bindings参数是一个列表顺序必须和engine中binding的顺序一致。如果你不确定顺序可以用engine.get_binding_name(i)逐个确认。我见过有人把输入和输出顺序搞反结果模型输出了一堆NaN。后处理YOLO特有的解码TensorRT的输出是原始的张量对于YOLOv8输出shape通常是(1, 84, 8400)其中84是4个坐标80个类别概率8400是预测框数量。你需要自己实现解码defpostprocess(output,conf_threshold0.5,iou_threshold0.45):# output shape: (1, 84, 8400)outputoutput.squeeze(0)# (84, 8400)outputoutput.transpose(1,0)# (8400, 84)boxes[]scores[]class_ids[]foriinrange(output.shape[0]):class_scoresoutput[i,4:]max_scoreclass_scores.max()ifmax_scoreconf_threshold:continueclass_idclass_scores.argmax()# 解码坐标这里假设是cx,cy,w,h格式cx,cy,w,houtput[i,:4]x1cx-w/2y1cy-h/2x2cxw/2y2cyh/2boxes.append([x1,y1,x2,y2])scores.append(max_score)class_ids.append(class_id)# NMSindicescv2.dnn.NMSBoxes(boxes,scores,conf_threshold,iou_threshold)# 处理结果...注意这里的坐标是相对于输入图像尺寸的需要缩放到原图尺寸。很多人忘记这一步导致检测框位置完全不对。序列化和反序列化的最佳实践在实际部署中engine的构建和推理通常是分开的。构建一次engine保存为文件后续推理时直接加载。但这里有个坑engine文件依赖具体的GPU型号和驱动版本。如果你在开发机上构建engine然后拷贝到生产服务器很可能加载失败。我的做法是在部署脚本中先检查是否存在engine文件如果不存在则构建存在则直接加载。同时在engine文件名中加入模型名称和输入尺寸的哈希值避免混淆。importhashlibdefget_engine_path(model_name,input_shape):hash_strf{model_name}_{input_shape}hash_valhashlib.md5(hash_str.encode()).hexdigest()[:8]returnf{model_name}_{hash_val}.engine另外序列化后的engine文件是二进制格式不要用文本编辑器打开也不要在不同字节序的机器间传输。我见过有人把engine文件用git管理结果每次pull下来都损坏就是因为git自动转换了换行符。性能调优别让GPU闲着如果你发现推理速度不达标可以从几个方向排查使用FP16或INT8精度在config中设置config.set_flag(trt.BuilderFlag.FP16)对于YOLO模型FP16几乎不影响精度但速度能提升30%-50%。INT8需要校准数据集麻烦一些但速度提升更明显。减少CPU-GPU数据传输如果输入数据是连续帧可以复用GPU内存避免每次推理都重新分配。我通常会在初始化时分配好内存推理时只做memcpy。使用多流并行如果你的应用需要同时处理多个视频流可以为每个流创建独立的Context和CUDA流实现并行推理。注意Context不是线程安全的每个线程必须有自己的Context。检查是否被CPU瓶颈拖累有时候预处理和后处理比推理本身还慢。考虑用GPU做预处理比如用CUDA核函数做resize和归一化。个人经验总结做了这么多年YOLO部署踩过的坑能写一本书。最后给几个实在的建议永远不要相信ONNX导出后的模型。导出后先用ONNX Runtime跑一遍确认结果正确再交给TensorRT。很多问题出在ONNX导出这一步TensorRT只是背锅侠。构建engine时加--verbose日志。虽然慢一点但能看清每一步优化做了什么出了问题好定位。保存engine文件时记录构建参数。我习惯在文件名里包含模型名、输入尺寸、精度模式、TensorRT版本号。这样即使过了半年看到文件名就知道这个engine是怎么来的。别在生产环境用动态batch。除非你真的需要否则静态batch省心省力。动态batch的构建时间、内存占用、推理延迟都不可控。测试时用真实数据。不要用全零或随机数据测试推理流程YOLO对输入数据很敏感全零输入可能输出正常但真实图片就崩了。用一张真实的测试图片从预处理到后处理全流程跑通。最后也是最重要的做好异常处理。TensorRT的API调用失败时有时候不会抛异常而是默默返回空指针。每个API调用后都要检查返回值别偷懒。这套流程我用了三年从YOLOv5到YOLOv8从Jetson Nano到A100基本没出过大问题。希望你能少走弯路。
