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

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

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

一、前置环境准备

1. 核心依赖安装

# 1. 基础依赖
pip install ultralytics==8.2.0 onnx==1.15.0 onnxsim==0.4.33 pycuda==2024.1

# 2. TensorRT安装（关键）
# 方式1：NVIDIA官方包（推荐，适配GPU驱动）
# 下载地址：https://developer.nvidia.com/tensorrt-download
tar -xzf TensorRT-8.6.1.6.Linux.x86_64-gnu.cuda-11.8.tar.gz
cd TensorRT-8.6.1.6/python
pip install tensorrt-8.6.1.6-cp39-none-linux_x86_64.whl

# 方式2：conda安装（仅适配部分版本）
conda install -c nvidia tensorrt==8.6.1

2. 环境验证

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

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

1. 训练/加载YOLOv8模型

from ultralytics import YOLO

# 加载预训练模型（或自有训练好的模型）
model = YOLO("yolov8s.pt")  # 替换为yolov8n.pt/yolov8l.pt，或自有权重yolov8s-best.pt

# 导出ONNX（固定输入尺寸，量化更稳定）
model.export(
    format="onnx",
    imgsz=640,  # 与训练时一致，边缘端可设为320/416
    batch=1,    # 量化仅支持batch=1
    simplify=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训练一致）。

import cv2
import numpy as np
import os

# 校准集路径（替换为你的校准集文件夹）
CALIB_IMG_DIR = "calib_images/"
CALIB_BATCH_SIZE = 1
CALIB_BATCH_NUM = 100  # 校准批次数量（建议≥100）

# 校准集加载函数（TensorRT要求的迭代器格式）
def load_calib_data():
    img_paths = [os.path.join(CALIB_IMG_DIR, f) for f in os.listdir(CALIB_IMG_DIR) if f.endswith((".jpg", ".png"))]
    for i in range(CALIB_BATCH_NUM):
        imgs = []
        for j in range(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))
        yield np.array(imgs)

2. 构建TensorRT量化引擎

import tensorrt as trt

# 1. 配置TensorRT logger
TRT_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_BATCH
parser = trt.OnnxParser(network, TRT_LOGGER)
with open("yolov8s.onnx", "rb") as f:
    parser.parse(f.read())

# 5. 构建并保存量化引擎
engine = builder.build_engine(network, config)
with open("yolov8s_int8.engine", "wb") as f:
    f.write(engine.serialize())
print("✅ INT8量化引擎保存完成：yolov8s_int8.engine")

四、量化模型推理验证

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

import tensorrt as trt
import cv2
import numpy as np

class YOLOv8TRTInfer:
    def __init__(self, engine_path):
        self.logger = trt.Logger(trt.Logger.ERROR)
        self.runtime = trt.Runtime(self.logger)
        with open(engine_path, "rb") as f:
            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")

    def infer(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.0
        img = np.expand_dims(img, 0).astype(np.float32)

        # 2. 分配显存（pycuda）
        import pycuda.driver as cuda
        import pycuda.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)
        for i in indices:
            i = i if isinstance(i, int) else i[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)
        return img_ori

# 推理示例
inferer = YOLOv8TRTInfer("yolov8s_int8.engine")
result_img = inferer.infer("test.jpg")
cv2.imwrite("result_int8.jpg", result_img)
print("✅ 推理完成，结果保存为result_int8.jpg")

2. 精度/速度对比验证

指标	YOLOv8s（PyTorch FP32）	YOLOv8s（TensorRT INT8）	变化率
mAP@0.5:0.95	44.9	42.5（典型值）	-5.3%（可接受）
FPS（V100）	110	280	+154%
模型体积	22MB（pt）	18MB（engine）	-18%

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

若量化后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频率。

---

常见问题速查

报错信息	解决方案
Unsupported operator xx	降低ONNX opset版本（如12→11），或用onnxsim简化
量化后精度暴跌	重新准备校准集，确保覆盖业务核心目标
推理时显存不足	降低max_workspace_size，或换更小的模型（yolov8s→yolov8n）