4. YOLOv5模型INT8量化与校准
任务目标
本案例使用TensorRT对YOLOv5s模型执行INT8量化，核心工作包括：构建校准数据集、实现校准器类、生成INT8引擎，并对比FP32与INT8引擎的推理延迟和精度差异。最终交付一个可直接运行的量化脚本，输出量化后的引擎文件及性能对比结果。
环境依赖
- TensorRT 8.x+（含Python API）
- PyTorch 1.10+（仅用于加载YOLOv5权重并转ONNX）
- OpenCV、NumPy
- YOLOv5官方仓库（提供模型定义与权重）
实现步骤
- 将YOLOv5s.pt转为FP32 ONNX（动态batch，固定输入640×640）
- 准备200张代表性图片作为校准集（建议来自目标检测训练集或真实场景）
- 实现Python校准器类，继承
IInt8EntropyCalibrator2，逐batch读取图片并预处理 - 使用TensorRT Builder构建INT8引擎，设置
set_flag(trt.BuilderFlag.INT8)并传入校准器 - 序列化引擎，加载后分别用FP32和INT8引擎推理同一批测试图，统计延迟与mAP
关键完整代码
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
import cv2
import os
import time
from glob import glob
# ---------- 校准数据加载器 ----------
class YOLOCalibrator(trt.IInt8EntropyCalibrator2):
def __init__(self, img_dir, batch_size=8, input_h=640, input_w=640):
super().__init__()
self.batch_size = batch_size
self.input_h = input_h
self.input_w = input_w
self.image_paths = glob(os.path.join(img_dir, '*.jpg'))[:200]
self.batch_idx = 0
self.device_input = cuda.mem_alloc(batch_size * 3 * input_h * input_w * np.float32().itemsize)
self.cache_file = 'yolov5s_int8.cache'
def get_batch_size(self):
return self.batch_size
def get_batch(self, names):
if self.batch_idx + self.batch_size > len(self.image_paths):
return None
batch_imgs = []
for i in range(self.batch_idx, self.batch_idx + self.batch_size):
img = cv2.imread(self.image_paths[i])
img = cv2.resize(img, (self.input_w, self.input_h))
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = img.astype(np.float32) / 255.0
img = img.transpose(2, 0, 1) # CHW
batch_imgs.append(img)
batch_data = np.ascontiguousarray(np.stack(batch_imgs), dtype=np.float32)
cuda.memcpy_htod(self.device_input, batch_data)
self.batch_idx += self.batch_size
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)
# ---------- 构建INT8引擎 ----------
def build_int8_engine(onnx_path, calib_dir, engine_path):
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(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(calib_dir)
config.int8_calibrator = calibrator
# 设置动态batch
profile = builder.create_optimization_profile()
input_tensor = network.get_input(0)
profile.set_shape(input_tensor.name, (1, 3, 640, 640), (8, 3, 640, 640), (16, 3, 640, 640))
config.add_optimization_profile(profile)
engine = builder.build_serialized_network(network, config)
with open(engine_path, 'wb') as f:
f.write(engine)
print(f'INT8 engine saved to {engine_path}')
# ---------- 推理与性能对比 ----------
def infer_engine(engine_path, test_img, warmup=10, repeats=50):
logger = trt.Logger(trt.Logger.WARNING)
runtime = trt.Runtime(logger)
with open(engine_path, 'rb') as f:
engine = runtime.deserialize_cuda_engine(f.read())
context = engine.create_execution_context()
context.set_input_shape('images', (1, 3, 640, 640))
img = cv2.imread(test_img)
img = cv2.resize(img, (640, 640))
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
img = img.transpose(2, 0, 1)[None] # NCHW
# 分配显存
d_input = cuda.mem_alloc(1 * 3 * 640 * 640 * np.float32().itemsize)
d_output = cuda.mem_alloc(1 * 25200 * 85 * np.float32().itemsize) # YOLOv5s输出
cuda.memcpy_htod(d_input, np.ascontiguousarray(img))
# 预热
for _ in range(warmup):
context.execute_v2([int(d_input), int(d_output)])
# 计时
times = []
for _ in range(repeats):
start = time.time()
context.execute_v2([int(d_input), int(d_output)])
cuda.Context.synchronize()
times.append((time.time() - start) * 1000)
return np.mean(times), np.std(times)
if __name__ == '__main__':
# 1. 转换ONNX（需先安装YOLOv5并下载权重）
# python export.py --weights yolov5s.pt --include onnx --dynamic
onnx_file = 'yolov5s.onnx'
calib_images = './calib_images' # 存放200张校准图片
int8_engine = 'yolov5s_int8.engine'
fp32_engine = 'yolov5s_fp32.engine' # 用trtexec或相同builder生成
build_int8_engine(onnx_file, calib_images, int8_engine)
# 2. 性能对比（需先生成FP32引擎）
test_image = './test.jpg'
int8_latency, int8_std = infer_engine(int8_engine, test_image)
fp32_latency, fp32_std = infer_engine(fp32_engine, test_image)
print(f'FP32 latency: {fp32_latency:.2f}±{fp32_std:.2f} ms')
print(f'INT8 latency: {int8_latency:.2f}±{int8_std:.2f} ms')
print(f'Speedup: {fp32_latency/int8_latency:.2f}x')运行与验证
- 准备校准集：创建
calib_images目录，放入200张含目标的真实图片（建议来自COCO或自采数据） - 生成FP32引擎：使用TensorRT自带的
trtexec工具：trtexec --onnx=yolov5s.onnx --saveEngine=yolov5s_fp32.engine - 执行量化：运行上述Python脚本，首次执行会生成
yolov5s_int8.cache校准缓存文件 - 验证通过条件：
- 脚本正常输出INT8引擎保存路径，无报错
- INT8引擎推理延迟显著低于FP32（通常快1.5～3倍）
- 用同一测试集分别跑FP32与INT8引擎的检测结果，mAP下降不超过3%（需额外实现NMS与mAP计算，此处不展开）
常见问题
- 校准集数量不足：少于100张可能导致量化后精度骤降，建议200张以上且覆盖各类目标
- 动态batch冲突：校准器batch_size需与profile中最小batch一致，否则报错
- 缓存复用：校准缓存文件生成后可复用，下次构建直接读取，无需重新校准