3. YOLOv8n模型INT8量化校准

任务说明

本案例使用TensorRT对YOLOv8n模型执行INT8量化校准，核心目标是：将FP32的YOLOv8n模型转换为INT8精度模型，并通过校准数据集统计激活值分布以确定量化尺度，最终对比INT8模型与FP32模型在验证集上的mAP指标，量化精度损失。

环境依赖

  • TensorRT ≥ 8.6（需支持INT8）
  • PyTorch ≥ 2.0 + torchvision（用于导出ONNX）
  • ultralytics（YOLOv8官方库，用于加载模型与数据集）
  • OpenCV、NumPy

实现原理与关键步骤

INT8量化校准的本质是确定每个张量的动态范围（scale）。TensorRT在校准阶段运行一组代表性输入（校准集），统计每层激活值的分布，使用直方图+熵校准（默认）或最小化KL散度的方法选择最优阈值，从而将FP32数值映射到INT8范围。本案例使用TensorRT Python API的INT8_CALIBRATOR接口，自定义校准器从COCO验证集中抽取子集完成校准。

关键完整代码

import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
import cv2
from ultralytics import YOLO
import os, time

# ---------- 1. 导出ONNX（若已存在可跳过） ----------
model = YOLO("yolov8n.pt")
model.export(format="onnx", opset=17, dynamic=False, imgsz=640)
onnx_path = "yolov8n.onnx"

# ---------- 2. 自定义INT8校准器 ----------
CALIB_CACHE = "yolov8n_calib.cache"
BATCH_SIZE = 8
CALIB_IMAGES = 128  # 校准图片数量

class YOLOCalibrator(trt.IInt8EntropyCalibrator2):
    def __init__(self, img_dir, cache_file):
        trt.IInt8EntropyCalibrator2.__init__(self)
        self.cache_file = cache_file
        self.image_paths = [os.path.join(img_dir, f) for f in os.listdir(img_dir)][:CALIB_IMAGES]
        self.batch_idx = 0
        self.device_input = cuda.mem_alloc(BATCH_SIZE * 3 * 640 * 640 * 4)
        self.shape = (BATCH_SIZE, 3, 640, 640)

    def get_batch_size(self):
        return BATCH_SIZE

    def get_batch(self, names):
        if self.batch_idx + BATCH_SIZE > len(self.image_paths):
            return None
        batch = []
        for i in range(BATCH_SIZE):
            img = cv2.imread(self.image_paths[self.batch_idx + i])
            img = cv2.resize(img, (640, 640))
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            img = img.astype(np.float32) / 255.0
            img = np.transpose(img, (2, 0, 1))  # CHW
            batch.append(img)
        self.batch_idx += BATCH_SIZE
        batch_np = np.ascontiguousarray(np.stack(batch), dtype=np.float32)
        cuda.memcpy_htod(self.device_input, batch_np)
        return [self.device_input]

    def read_calibration_cache(self):
        if os.path.exists(self.cache_file):
            return open(self.cache_file, "rb").read()
        return None

    def write_calibration_cache(self, cache):
        with open(self.cache_file, "wb") as f:
            f.write(cache)

# ---------- 3. 构建INT8引擎 ----------
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(TRT_LOGGER)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, TRT_LOGGER)
with open(onnx_path, "rb") as f:
    parser.parse(f.read())

config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB
config.set_flag(trt.BuilderFlag.INT8)
calibrator = YOLOCalibrator("coco_val_images/", CALIB_CACHE)
config.int8_calibrator = calibrator

engine = builder.build_serialized_network(network, config)
with open("yolov8n_int8.engine", "wb") as f:
    f.write(engine)
print("INT8引擎构建完成")

# ---------- 4. 精度对比评估 ----------
def run_inference(engine_path, image):
    runtime = trt.Runtime(TRT_LOGGER)
    with open(engine_path, "rb") as f:
        engine = runtime.deserialize_cuda_engine(f.read())
    context = engine.create_execution_context()
    # 预处理
    img = cv2.resize(image, (640, 640))
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
    img = np.transpose(img, (2, 0, 1))[None]
    # 分配内存
    output_shape = (1, 84, 8400)
    d_input = cuda.mem_alloc(img.nbytes)
    d_output = cuda.mem_alloc(np.prod(output_shape) * 4)
    cuda.memcpy_htod(d_input, np.ascontiguousarray(img))
    context.execute_v2([int(d_input), int(d_output)])
    output = np.empty(output_shape, dtype=np.float32)
    cuda.memcpy_dtoh(output, d_output)
    return output

# 使用COCO验证集子集评估mAP（此处仅演示单张推理流程）
val_img = cv2.imread("coco_val_images/000000000139.jpg")
out_fp32 = run_inference("yolov8n_fp32.engine", val_img)  # 需先构建FP32引擎
out_int8 = run_inference("yolov8n_int8.engine", val_img)
print("FP32输出shape:", out_fp32.shape)
print("INT8输出shape:", out_int8.shape)
# 完整mAP评估请使用ultralytics的Val API，加载engine文件进行批量验证

运行与验证

  1. 准备校准数据：从COCO val2017中随机抽取128张图片放入coco_val_images/目录（需先下载并解压）。
  2. 构建FP32引擎（用于对比）：将上述代码中INT8标志移除，构建yolov8n_fp32.engine。
  3. 执行量化：运行上述脚本，TensorRT会读取校准图片并生成yolov8n_calib.cache缓存文件，最终输出INT8引擎。
  4. 精度评估：使用ultralytics库加载两个引擎，在COCO val2017上分别运行验证，对比mAP@0.5:0.95指标。

通过条件：INT8引擎成功构建且能正常推理；INT8模型mAP相对FP32下降不超过3%（典型值1~2%）；单张640×640推理延迟较FP32降低30%以上（实际值取决于GPU型号）。

常见问题

  • 校准缓存未生成：检查校准图片路径是否正确，以及get_batch返回None的终止条件是否触发。
  • 精度下降过大：增加校准图片数量至256~512张，或改用IInt8MinMaxCalibrator（适合检测任务）。
  • 显存不足：降低BATCH_SIZE至4或2，或减小workspace内存池限制。
← 上一章 📖 返回目录 下一章 →