ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

YOLOv5 ONNX模型TensorRT INT8量化实战:从标定到多路部署

YOLOv5 ONNX模型TensorRT INT8量化实战:从标定到多路部署 简介面向嵌入式、边缘计算等资源受限场景中的深度学习开发者资源演示了基于 Python 使用 TensorRT 对 YOLOv5 ONNX 模型进行 INT8 量化的完整实现。INT8 量化将模型计算从 FP32 浮点转换为 8 位整数可有效降低显存占用和推理延迟过程中包含校准与构建两个关键步骤校准数据的代表性直接影响量化后模型的精度上限。RAR 压缩包共 10 个文件包含 .py 源码与 .pyc 编译字节码、.trt 引擎文件以及 .cache 校准缓存文件包体大小仅 6.95MB便于离线分析和快速复现。资源内提供校准器、转换量化脚本、工具模块等核心代码并附有 yolov5s 与 nanodet 的量化后引擎和校准缓存可以对照不同模型的量化效果调整参数并验证精度。已有 4473 人学习下载适合希望掌握模型压缩与推理加速、在资源受限设备上部署实时目标检测应用的开发者参考。1. INT8 量化 YOLOv5用激活分布换吞吐不是拿精度赌博手头只有一张 T4却要同时处理几路 1080p 25fps 视频流这种时候最憋屈的不是模型跑不动而是 GPU 显存带宽先见底。FP16 在大多数 GPU 上算力都够真正的瓶颈是每帧要搬运的张量宽度把 YOLOv5 的 ONNX 模型交给 TensorRT 做 INT8 量化等于把大部分中间张量从 2 字节压到 1 字节访存压力直接砍掉一半。这篇笔记围绕「基于 Python 的 TensorRT INT8 量化落地到 YOLOv5 ONNX 模型」展开从导出标准 ONNX 开始到手写标定器再到 Python 侧推理与多路并发最后是几个我反复踩的坑。适合正在做边缘盒子、GPU 服务或视频结构化部署的工程师新手能照着跑熟手可以直接看参数和边界。2. 先有标准 ONNX再有 INT8opset、动态轴和 FP16 基线TensorRT 不能直接消费 PyTorch 的 .pt 文件它的入口是 ONNX 或者用 TensorRT API 自己搭图。所以流程的第一步是把 YOLOv5 导出成一个干净、可复现的 ONNX。这一步不直接决定最终精度但决定后面标定能不能顺利跑通。我一般会把导出、结构校验、FP16 引擎构建这三件事放在同一天做完后面 INT8 出问题时才有对照基线。2.1 导出 ONNX 时把三件事锁死opset、simplify、动态轴官方仓库自带 export.py不建议自己拼torch.onnx.export的参数版本匹配问题太多。我用的命令python export.py \ --weights yolov5s.pt \ --include onnx \ --opset 17 \ --dynamic \ --simplify这里三个参数是关键--opset 17保证 ONNX 算子版本足够新TensorRT 对Pow、Sqrt、Resize这类算子的解析在较新 opset 下更完整--simplify会调用 onnx-simplifier 做常量折叠把 Focus 或早期 Conv 结构里的切片、拼接化简避免后续 trtexec 报莫名其妙的形状推导错误--dynamic导出动态 batch 和动态宽高后面用 trtexec 构建时再通过 min/opt/max shapes 重新圈定范围。导出后马上做结构校验。这一步能拦住一半的「构建失败」问题import onnx model onnx.load(yolov5s.onnx) onnx.checker.check_model(model) print([inp.name for inp in model.graph.input]) # 预期只有一个输入images print([out.name for out in model.graph.output]) # 预期只有一个输出output0检查点有两个输入只有一个名字通常是images输出也只有一个YOLOv5 导出时会把三个检测头的原始预测拼成一个[batch, 8400, 5nc]张量后处理不留在 ONNX 图里。如果你在图上看到了NonMaxSuppression算子说明导出配置带了额外后处理建议去掉不然 TensorRT 的层融合会被 NMS 拖住INT8 校准也容易因为输出尺度异常而失败。2.2 先跑一版 FP16 引擎当对照组trtexec 两行命令不要一上来就做 INT8。先构建一个 FP16 引擎把延迟和精度的基线测出来后面 INT8 的一切结论才有参照物。构建命令trtexec \ --onnxyolov5s.onnx \ --saveEngineyolov5s_fp16.engine \ --fp16 \ --minShapesimages:1x3x640x640 \ --optShapesimages:1x3x640x640 \ --maxShapesimages:1x3x640x640 \ --workspace4096参数说明--fp16对应构建时打开 FP16 开关--minShapes/--optShapes/--maxShapes把输入形状的范围固定成同一个值等于关闭动态 shape这是跑通流程最省事的方式。--workspace4096是旧版参数新版本 TensorRT 改成了--memPoolSizeworkspace:4096单位是 MiB含义是允许 TensorRT 用多少显存做层融合和格式选择建议 2GB 到 4GB太小会导致部分算子不能融合。注意T4 这类卡不支持 FP8FP8 是 Hopper 及之后架构才有的东西。在 T4 上做量化对比只在 FP16 和 INT8 之间比就够了别浪费时间纠结 FP8。3. 写一个 YOLOv5 专用的 INT8 校准器标定集、缓存和不均匀分布FP16 是直接把 float 的指数位砍半精度损失是均匀可预测的INT8 是均匀分箱但每一层的激活分布形状差很多YOLOv5 前几层输出动态范围特别大后层小目标对应的激活又很小。所以需要一个标定集去逐层统计激活分布再为每层挑一个合适的量化阈值。这一步做不好mAP 掉 5 个点都是轻的。3.1 选 Calibrator 先看这张对照表Entropy、MinMax、LegacyTensorRT 的 Python API 里常用三种标定器它们的差异在阈值选择策略上标定器阈值策略适合场景主要风险IInt8EntropyCalibrator2最小化量化前后分布的 KL 散度目标检测、分类的通用默认选择标定集类别不均衡时小目标阈值会偏移IInt8MinMaxCalibrator直接取激活的 min/max 作为阈值激活分布集中在极值范围内时对离群点敏感整体噪声变大IInt8LegacyCalibrator早期熵标定的实现兼容旧脚本效果不优于 Calibrator2不建议新项目用对 YOLOv5 这种多尺度检测模型我默认用IInt8EntropyCalibrator2。它的核心逻辑是试探不同的截断阈值看哪个阈值让量化前后的分布差异最小。MinMax 在面对 640×640 输入产生的那些极端激活值时很容易把阈值拉得过大中间段的激活全被压到相邻的量化等级里小目标就丢了。3.2 标定集不是随机截图类别均衡、光暗变化和预处理一致标定集是整个 INT8 流程里最容易被低估的环节。常见错误是随便抓 100 张视频帧丢进目录结果类别失衡白天场景占了 90%夜间目标全部消失。我一般按三个原则组织标定集。类别均衡每个类别至少 5 张长尾类别可以多放标定集够用就行50 到 200 张之间多了反而把分布抹平少了覆盖不了极端场景。光暗变化把白天、黄昏、夜间、逆光、低照度按比例混进去尤其是夜间场景YOLOv5 在暗光下的激活分布和白天天差地别标定集里没有暗光样本INT8 模型一到晚上就翻车。预处理一致标定图片在喂给标定器之前必须走和线上推理完全相同的 letterbox、BGR 转 RGB、除以 255 这三步。很多人标定时用普通 resize 到 640×640推理时用 letterboxINT8 的量化阈值完全错位。3.3 Python 实现 EntropyCalibrator2 的最小类下面这个类可以直接用。前置条件是先把标定图片经过预处理存成 NCHW 的 float32 数组值域 0 到 1每张图一个 .npy 文件。import glob import os import numpy as np import tensorrt as trt class YOLOv5Calibrator(trt.IInt8EntropyCalibrator2): def __init__(self, calib_dir, batch_size, cache_file): super().__init__() self.calib_dir calib_dir self.batch_size batch_size self.cache_file cache_file self.files sorted(glob.glob(os.path.join(calib_dir, *.npy))) self.idx 0 self._batch None # 保持引用防止 numpy 内存被回收 def get_batch_size(self): return self.batch_size def get_batch(self, names): if self.idx self.batch_size len(self.files): return None # 返回 None 表示标定数据取完 batch_paths self.files[self.idx : self.idx self.batch_size] batch np.stack([np.load(p) for p in batch_paths]).astype(np.float32) self.idx self.batch_size self._batch batch # YOLOv5 只有一个输入names 里只有一个名字给每个名字同一个地址 return [int(self._batch.ctypes.data) for _ in names] def read_calibration_cache(self): # 第二次构建时直接读 cache跳过标定 if os.path.exists(self.cache_file): with open(self.cache_file, rb) as f: return f.read() return None def write_calibration_cache(self, cache): # 先写临时文件再改名避免程序中断留下半截 cache tmp self.cache_file .tmp with open(tmp, wb) as f: f.write(cache) os.replace(tmp, self.cache_file)参数说明get_batch返回的是一个指针列表TensorRT 标定流程会从这块 CPU 内存里读数据再拷进 GPU。所以这里特别强调把self._batch挂住如果只写局部变量函数返回后 numpy 的内存可能被回收TRT 读到的就是野指针典型症状是构建过程中随机崩溃或者标定结果每次都不一样。read_calibration_cache和write_calibration_cache配合使用第一次标定要跑完整的数据集第二次构建直接读 cache几秒钟就能出引擎。3.4 用构建器把校准器接上去落盘 INT8 引擎标定器写好后剩下的就是把 INT8 开关打开构建引擎logger trt.Logger(trt.Logger.WARNING) builder trt.Builder(logger) network builder.create_network(1 int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) parser trt.OnnxParser(network, logger) with open(yolov5s.onnx, rb) as f: parser.parse(f.read()) config builder.create_builder_config() config.set_flag(trt.BuilderFlag.INT8) config.int8_calibrator YOLOv5Calibrator( calib_dircalib_npy, batch_size8, cache_fileyolov5s_int8.cache, ) # 这行看情况开。开了之后 TRT 会保留一部分 FP16 算子 config.set_flag(trt.BuilderFlag.FP16) config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 30) engine_bytes builder.build_serialized_network(network, config) with open(yolov5s_int8.engine, wb) as f: f.write(engine_bytes)这段代码里有个常见混淆点config.set_flag(trt.BuilderFlag.FP16)和 INT8 同时打开不代表模型是纯 INT8。TensorRT 会保留那些量化误差大的算子走 FP16这是合理的工程选择。如果你希望严格验证纯 INT8 效果就把 FP16 开关去掉重新构建一次。另外注意config.int8_calibrator calibrator是较新 API 的写法老版本用的是config.set_calibrator(calibrator)两者行为一致但别混用。提示换 GPU 型号、换 TensorRT 版本或者换标定集之后cache 文件必须删除重建。这一点后面避坑章节还会展开。4. 在 Python 里把 INT8 引擎用起来输入绑定、后处理与多路并发engine 文件不是模型文件是一个针对「GPU 型号 TensorRT 版本 输入 shape」序列化出来的二进制产物。它不能跨卡复用换卡必须重新构建。Python 侧的工作核心是反序列化引擎、绑定输入输出内存、用和标定一致的预处理喂数据、把输出张量后处理成坐标框。4.1 反序列化绑定内存只分配一次循环里不要 newimport numpy as np import tensorrt as trt import pycuda.driver as cuda import pycuda.autoinit class TRTEngine: def __init__(self, engine_path): logger trt.Logger(trt.Logger.WARNING) with open(engine_path, rb) as f, trt.Runtime(logger) as runtime: self.engine runtime.deserialize_cuda_engine(f.read()) self.context self.engine.create_execution_context() self.stream cuda.Stream() self.inputs [] self.outputs [] self.bindings [] for i in range(self.engine.num_bindings): name self.engine.get_binding_name(i) shape self.engine.get_binding_shape(i) dtype trt.nptype(self.engine.get_binding_dtype(i)) host cuda.pagelocked_empty(trt.volume(shape), dtype) device cuda.mem_alloc(host.nbytes) binding {name: name, shape: shape, dtype: dtype, host: host, device: device} if self.engine.binding_is_input(i): self.inputs.append(binding) else: self.outputs.append(binding) self.bindings.append(device)这里用cuda.pagelocked_empty而不是普通 numpy 数组是因为页锁定内存的地址稳定D2H 拷贝走 DMA 更快多路并发时不容易出现地址被系统换页的情况。绑定动作必须在循环外面完成常见翻车写法是在每帧推理里重新cuda.mem_alloc几小时跑下来显存碎片化然后就 OOM 了。如果引擎是动态 shape绑定时要先调用context.set_binding_shape(0, actual_shape)再根据实际体积分配 buffer。我这里 mon 不上动态 batch先用静态 shape 把流程跑通。4.2 推理与 letterbox 一一对应预处理错一点量化优势全还回去def letterbox_to_640(img, target640): h, w img.shape[:2] ratio min(target / w, target / h) new_w, new_h int(round(w * ratio)), int(round(h * ratio)) resized cv2.resize(img, (new_w, new_h), interpolationcv2.INTER_LINEAR) dw, dh (target - new_w) // 2, (target - new_h) // 2 canvas np.full((target, target, 3), 114, dtypenp.uint8) canvas[dh:dh new_h, dw:dw new_w] resized # BGR - RGB、HWC - CHW、归一化到 0-1 tensor canvas[:, :, ::-1].transpose(2, 0, 1).astype(np.float32) / 255.0 return np.ascontiguousarray(tensor), ratio, dw, dh这里每一行都有讲究。dw, dh是 letterbox 时补边的偏移量后处理阶段把输出坐标映射回原图时要用到。填充值 114 是 YOLOv5 训练时默认的 pad 值不要随手改成 0。BGR 转 RGB 用canvas[:, :, ::-1]实现YOLOv5 训练时用的是 RGB 输入。最后除以 255这一步千万不能忘标定数据如果也是按这个流程生成的两边就对齐了。推理循环里的拷贝和同步动作也要固定def infer(self, input_tensor): np.copyto(self.inputs[0][host], input_tensor.ravel()) cuda.memcpy_htod_async(self.inputs[0][device], self.inputs[0][host], self.stream) self.context.execute_async_v2(bindingsself.bindings, stream_handleself.stream.handle) cuda.memcpy_dtoh_async(self.outputs[0][host], self.outputs[0][device], self.stream) cuda.stream.synchronize(self.stream) return self.outputs[0][host]stream.synchronize()是必须的不然下一帧np.copyto写入 host 内存时上一帧的 D2H 拷贝可能还没完成读出来的是新旧数据混合。在多路场景里可以让每个路经持有自己的 cuda.Stream互不阻塞。4.3 后处理把输出张量从 8400×85 收缩到最终框YOLOv5 在 640 输入下的原始输出是[1, 8400, 85]。8400 来自三个尺度的锚点网格80×80 检测小目标40×40 检测中目标20×20 检测大目标合计 840085 的前 4 个是 cx、cy、w、h第 5 个是 objectness 置信度后面 80 个是 COCO 类别分数。如果是自定义数据集形状变成[1, 8400, 5nc]代码用nc pred.shape[-1] - 5自动适配。import torch from torchvision.ops import nms def postprocess(pred, conf_thres0.25, iou_thres0.45, ratio1.0, dw0, dh0): pred torch.from_numpy(pred) nc pred.shape[-1] - 5 scores pred[..., 4] * pred[..., 5:].max(dim-1).values xy pred[..., :2] wh pred[..., 2:4] x1y1 xy - wh / 2 x2y2 xy wh / 2 xyxy torch.cat([x1y1, x2y2], dim-1) mask scores conf_thres if not mask.any(): return [], [], [] boxes xyxy[mask] confs scores[mask] keep nms(boxes, confs, iou_thres) boxes boxes[keep] / ratio boxes[:, [0, 2]] - dw boxes[:, [1, 3]] - dh return boxes, confs[keep]这里有个容易被忽略的细节torchvision.ops.nms是类别无关的全局 NMS不同类别的两个框如果高度重叠后出现的会被压掉。在密集场景下比如两个人挨着站这种策略偶尔会丢掉其中一个。要更严谨就按类别分组做 NMS遍历每个类别的索引分别调用一次 nms代码量不大但能减少一类误伤。4.4 多路 1080p 视频流一个引擎多 context并发而不是翻倍 batch很多人在 T4 上做 25fps 视频流检测时第一反应是把 batch 调到 4 或 8。对大模型这是对的但 YOLOv5s 在 640 下本身算子不重大 batch 反而容易让后处理和 NMS 卡在 CPU 侧。常见做法是 batch1开多路并发一个 engine 只反序列化一次每路各创建一个 context权重在显存里只有一份。engine TRTEngine(yolov5s_int8.engine) workers [Worker(engine) for _ in range(N)]每个 Worker 内部独立持有 context、输入输出 buffer 和 cuda.Stream推理循环完全独立。这个方法能撑多少路不要凭感觉按下面几步实测第一步用 trtexec 分别测 FP16 和 INT8 引擎在 640×640 下的平均 GPU 延迟记作 L_fp16 和 L_int8。第二步用 Python 侧完整推理链单独测一帧的预处理加后处理 CPU 耗时。第三步按公式估算GPU 侧理论上限是1000 / (L_gpu × fps)路CPU 侧理论上限是1000 / (T_cpu × fps)路取两者较小值。比如 INT8 GPU 延迟 3ms25fps 需求下单路每秒占 GPU 75ms理论上限约 13 路如果 CPU 预处理加 NMS 每帧要 4ms单路每秒 100ms那 CPU 先到瓶颈上限是 10 路这时候该优化的是后处理不是模型。5. INT8 量化避坑记录5 个翻车场景和我的处理顺序INT8 的坑不像 FP16 那样直观很多问题看起来是模型精度不行实际上出在标定流程或工程实现上。这里列 5 个我遇到过的场景每条按现象、原因、解决展开。5.1 换 GPU 后复用 cache 文件框直接漂移现象在 A 卡上标定生成的 cache 文件拿到 B 卡上复用INT8 引擎跑同一张图置信度普遍下降 0.1 以上边缘框位置轻微晃动。FP16 引擎没有这个问题。原因校准缓存不只是一个数学分布文件它和具体 GPU 的算子实现、卷积融合方式有耦合。换卡之后中间层激活分布的实际数值边界变了但 cache 里的统计阈值还是旧卡的量化误差被放大。解决cache 文件的命名里带上 GPU 型号和 TensorRT 版本例如yolov5s_t4_trt8601.cache。构建脚本里判断这个文件不存在就直接删掉重新标定不要留一个模棱两可的通用名字。5.2 标定图片没走同一套 letterbox小目标全灭现象标定时用 cv2 直接把图片 resize 到 640×640推理时用 letterboxINT8 模型在远处小目标上几乎全部丢失FP16 正常。原因普通 resize 改变了图像宽高比远处目标被压扁模型输出的激活分布和线上推理完全不同。标定器统计出来的阈值套不到推理时的真实激活上。解决把 letterbox、BGR 转 RGB、归一化抽成一个公共函数标定脚本和推理脚本都 import 同一个函数。这个错误最隐蔽的地方是单看 mAP 可能没掉多少但小目标确实没了因为小目标对应的高频激活只占分布尾部的一小部分。5.3 量化后置信度普遍偏低但框的位置是对的现象检测框位置基本准确但置信度从 FP16 的 0.9 掉到 0.3后处理阈值 0.25 直接过滤掉一半。原因YOLOv5 的 objectness 输出在 INT8 量化后落在了一个很窄的量化区间里步长限制了置信度分辨率但不影响坐标回归。这是输出层量化噪声的表现不一定是标定集的问题。解决先把 conf_thres 降到 0.15 看框是否还在。如果框还在就在后处理里对 objectness 做一个经验性补偿系数比如乘 1.1 或 1.2如果框也丢了再回到标定集找原因。不要一上来就重新标定浪费几个小时。5.4 多路并发时偶发 CUDA errorout of memory现象单路跑一整天没问题加到四路之后每隔几分钟报一次 CUDA OOM进程崩溃重启恢复。原因每一路都重新反序列化了一次 engine。TensorRT 反序列化之后权重在图里的显存副本是独立的四路就有四份权重。同时每路还各自分配 input 和 output 的 pinned memory累积起来很快把显存和内存都吃满。解决全局只反序列化一次 engine每个 worker 进程或线程里engine.create_execution_context()创建独立 context。context 共享权重只额外占用少量执行上下文空间。还要注意 context 不是线程安全的同一个 context 不能给多个线程并发调用。5.5 cache 写坏了之后没有快捷重标定的通道现象标定过程中程序被 killcache 文件写了一半。下次构建时 TensorRT 直接报错或者更隐蔽——不报错但精度突然变差。原因write_calibration_cache默认直接打开目标路径写文件进程中断会把 cache 截断。下次构建时 TRT 可能读到一段不完整的字节流。解决写临时文件再os.replace原子替换代码在 3.3 节已经给了。同时把重标定命令固化成一个脚本遇到精度异常时第一件事就是删 cache 重标定而不是去调后处理阈值。6. 顺手压一压延迟用 CUDA Event 校准对比让 cache 也有后悔药INT8 做完之后先别急着改后处理。我会用 CUDA Event 在 GPU 侧精确测量一段推理耗时不然 Python 的time.time()会把 D2H 拷贝和 CPU 启动开销全混进去数字没法指导后续调优。start cuda.Event() end cuda.Event() start.record(engine.stream) engine.context.execute_async_v2( bindingsengine.bindings, stream_handleengine.stream.handle, ) end.record(engine.stream) end.synchronize() print(gpu ms:, start.time_till(end))time.till返回的是毫秒级的 GPU 时间测量范围只覆盖execute_async_v2一次推理不包含 memcpy 和后处理。如果要完整测端到端延迟把 H2D、推理、D2H 都包进来分别测别混在一起。另一个习惯是给 cache 和 engine 都留后悔药。我每改动一次标定集、输入尺寸、TensorRT 版本就会把新生成的 cache 文件复制一份到带日期和改动说明的目录比如backup/yolov5s_int8_t4_20250112/。线上 INT8 出问题时第一反应不是重新标定而是拿上一版备份 cache 重新构建引擎对比两边的推理输出快速定位是改动引入的回归还是数据分布变化。这样能省掉很多重复标定时间。我的工作习惯是先量 FP16 基线再动 INT8出问题至少有一瓶后悔药可喝。这个顺序帮我避开过很多次「越调越乱」的局面。希望帮到你。本文还有配套的精品资源点击获取
返回列表