YOLOv8模型TensorRT量化实操步骤手册

YOLOv8模型TensorRT量化实操步骤手册

本手册聚焦YOLOv8（检测任务）的 TensorRT INT8 量化（工业界主流轻量化方案），覆盖「环境准备→模型导出→量化转换→推理验证→调优」全流程，适配 Linux（Ubuntu 20.04）/Windows 10+，硬件参考 NVIDIA GPU（RTX 3090/Tesla V100）。

一、前置环境准备

1. 核心依赖安装

# 1. 基础依赖pipinstallultralytics==8.2.0onnx==1.15.0onnxsim==0.4.33pycuda==2024.1# 2. TensorRT安装（关键）# 方式1：NVIDIA官方包（推荐，适配GPU驱动）# 下载地址：https://developer.nvidia.com/tensorrt-downloadtar-xzfTensorRT-8.6.1.6.Linux.x86_64-gnu.cuda-11.8.tar.gzcdTensorRT-8.6.1.6/python pipinstalltensorrt-8.6.1.6-cp39-none-linux_x86_64.whl# 方式2：conda安装（仅适配部分版本）condainstall-cnvidiatensorrt==8.6.1

2. 环境验证

importtensorrtastrtimportultralyticsprint(f"TensorRT版本：{trt.__version__}")# 需≥8.0print(f"YOLOv8版本：{ultralytics.__version__}")# 需≥8.0

二、YOLOv8模型导出ONNX（量化前置步骤）

1. 训练/加载YOLOv8模型

fromultralyticsimportYOLO# 加载预训练模型（或自有训练好的模型）model=YOLO("yolov8s.pt")# 替换为yolov8n.pt/yolov8l.pt，或自有权重yolov8s-best.pt# 导出ONNX（固定输入尺寸，量化更稳定）model.export(format="onnx",imgsz=640,# 与训练时一致，边缘端可设为320/416batch=1,# 量化仅支持batch=1simplify=True,# 简化ONNX算子，避免量化报错opset=12# TensorRT8.x推荐opset=12)# 导出后生成：yolov8s.onnx

2. ONNX模型验证

# 检查ONNX算子兼容性polygraphy inspect model yolov8s.onnx --show-inputs --show-outputs

✅ 正常输出：输入images(1,3,640,640)，输出output0(1,84,8400)。

三、TensorRT INT8量化实操（核心步骤）

1. 准备量化校准集（关键！影响量化精度）

• 要求：选取100-500张与测试集分布一致的图片（无需标注），避免纯黑/纯白图；
• 格式：转为RGB格式，尺寸640×640，归一化至0-1（与YOLOv8训练一致）。

importcv2importnumpyasnpimportos# 校准集路径（替换为你的校准集文件夹）CALIB_IMG_DIR="calib_images/"CALIB_BATCH_SIZE=1CALIB_BATCH_NUM=100# 校准批次数量（建议≥100）# 校准集加载函数（TensorRT要求的迭代器格式）defload_calib_data():img_paths=[os.path.join(CALIB_IMG_DIR,f)forfinos.listdir(CALIB_IMG_DIR)iff.endswith((".jpg",".png"))]foriinrange(CALIB_BATCH_NUM):imgs=[]forjinrange(CALIB_BATCH_SIZE):img=cv2.imread(img_paths[i*CALIB_BATCH_SIZE+j])img=cv2.cvtColor(img,cv2.COLOR_BGR2RGB)img=cv2.resize(img,(640,640))img=img.transpose(2,0,1)/255.0# 转为CHW，归一化imgs.append(img.astype(np.float32))yieldnp.array(imgs)

2. 构建TensorRT量化引擎

importtensorrtastrt# 1. 配置TensorRT loggerTRT_LOGGER=trt.Logger(trt.Logger.INFO)builder=trt.Builder(TRT_LOGGER)config=builder.create_builder_config()# 2. 设置量化参数config.set_flag(trt.BuilderFlag.INT8)# 开启INT8量化config.int8_calibrator=trt.IInt8EntropyCalibrator2(# 熵校准（主流方案）training_data=load_calib_data(),cache_file="yolov8s_calib.cache",# 校准缓存文件（避免重复校准）batch_size=CALIB_BATCH_SIZE)# 3. 设置显存限制（根据GPU调整）config.max_workspace_size=1<<30# 1GB，可设为1<<32（4GB）# 4. 解析ONNX模型并构建引擎network=builder.create_network(1<<0)# EXPLICIT_BATCHparser=trt.OnnxParser(network,TRT_LOGGER)withopen("yolov8s.onnx","rb")asf:parser.parse(f.read())# 5. 构建并保存量化引擎engine=builder.build_engine(network,config)withopen("yolov8s_int8.engine","wb")asf:f.write(engine.serialize())print("✅ INT8量化引擎保存完成：yolov8s_int8.engine")

四、量化模型推理验证

1. TensorRT推理代码（适配YOLOv8输出解析）

importtensorrtastrtimportcv2importnumpyasnpclassYOLOv8TRTInfer:def__init__(self,engine_path):self.logger=trt.Logger(trt.Logger.ERROR)self.runtime=trt.Runtime(self.logger)withopen(engine_path,"rb")asf:self.engine=self.runtime.deserialize_cuda_engine(f.read())self.context=self.engine.create_execution_context()# 获取输入/输出索引self.input_idx=self.engine.get_binding_index("images")self.output_idx=self.engine.get_binding_index("output0")definfer(self,img_path):# 1. 图片预处理（与校准集一致）img=cv2.imread(img_path)img_ori=img.copy()img=cv2.cvtColor(img,cv2.COLOR_BGR2RGB)img=cv2.resize(img,(640,640))img=img.transpose(2,0,1)/255.0img=np.expand_dims(img,0).astype(np.float32)# 2. 分配显存（pycuda）importpycuda.driverascudaimportpycuda.autoinit d_input=cuda.mem_alloc(img.nbytes)d_output=cuda.mem_alloc(1*84*8400*4)# float32=4字节# 3. 数据拷贝+推理cuda.memcpy_htod(d_input,img)bindings=[int(d_input),int(d_output)]self.context.execute_v2(bindings)# 4. 结果拷贝+解析output=np.empty((1,84,8400),dtype=np.float32)cuda.memcpy_dtoh(output,d_output)# 5. YOLOv8输出解析（简化版，可替换为ultralytics原生解析）output=output[0].reshape(-1,84)# (8400, 84)boxes=output[:,:4]# x1,y1,x2,y2（归一化）scores=output[:,4:].max(axis=1)cls_ids=output[:,4:].argmax(axis=1)# 6. 非极大值抑制（NMS）+ 绘制结果indices=cv2.dnn.NMSBoxes(boxes[:,:4].tolist(),scores.tolist(),0.25,0.45)foriinindices:i=iifisinstance(i,int)elsei[0]x1,y1,x2,y2=boxes[i]# 反归一化到原图尺寸h,w=img_ori.shape[:2]x1,y1=int(x1*w),int(y1*h)x2,y2=int(x2*w),int(y2*h)cv2.rectangle(img_ori,(x1,y1),(x2,y2),(0,255,0),2)cv2.putText(img_ori,f"cls{cls_ids[i]}:{scores[i]:.2f}",(x1,y1-10),cv2.FONT_HERSHEY_SIMPLEX,0.5,(0,255,0),2)returnimg_ori# 推理示例inferer=YOLOv8TRTInfer("yolov8s_int8.engine")result_img=inferer.infer("test.jpg")cv2.imwrite("result_int8.jpg",result_img)print("✅ 推理完成，结果保存为result_int8.jpg")

2. 精度/速度对比验证

五、量化精度下降调优（关键！）

若量化后mAP下降超过10%，按以下优先级调优：

1. 优化校准集：增加校准集数量（≥200张），确保覆盖所有目标类别/尺寸；
2. 量化感知训练（QAT）：在YOLOv8训练时插入量化节点（需修改ultralytics源码）；
3. 跳过敏感层量化：对特征融合层（FPN）、注意力层保留FP32精度；
4. 降低量化强度：改用FP16量化（精度几乎无损，速度提升~80%）：

   # FP16量化（替换INT8配置）config.set_flag(trt.BuilderFlag.FP16)# 无需校准集，直接构建引擎

六、边缘端部署适配

1. 导出为TensorRT引擎后：可直接部署到Jetson Nano/Xavier（需对应ARM版本TensorRT）；
2. 批量推理优化：若需批量处理，可修改batch_size为4/8（需重新导出ONNX+构建引擎）；
3. 低功耗模式：Jetson设备可通过jetson_clocks工具调优GPU频率。

常见问题速查