从YOLOv5 ONNX到TensorRT INT8引擎：一次量化实践与踩坑记录

1. YOLOv5模型INT8量化实战指南

最近在部署YOLOv5模型时，我发现INT8量化真是个好东西。它能显著减小模型体积，提升推理速度，特别适合边缘设备部署。不过实际操作起来，从ONNX到TensorRT INT8引擎的转换过程并不像想象中那么顺利。下面我就把整个实践过程，包括遇到的坑和解决方案，完整分享给大家。

先说说为什么要做INT8量化。简单来说，就是把模型参数从32位浮点（FP32）压缩到8位整数（INT8）。这样做的好处显而易见：模型体积能缩小4倍，推理速度也能提升2-3倍。但代价是精度会有轻微下降，不过在很多实际应用中，这点精度损失完全在可接受范围内。

我这次使用的环境是：

• Ubuntu 20.04
• CUDA 11.3
• TensorRT 8.2
• PyTorch 1.10
• YOLOv5 v6.1版本

2. 准备工作与环境搭建

2.1 模型导出为ONNX格式

首先需要把训练好的YOLOv5模型导出为ONNX格式。这里有个小技巧，导出时建议加上dynamic参数，这样生成的ONNX模型可以支持动态batch size：

python复制python export.py --weights yolov5s.pt --include onnx --dynamic

导出时可能会遇到一些警告，只要不是错误就可以忽略。特别注意要记下模型的输入尺寸，后续量化时会用到。我这次用的是640x640的输入尺寸。

2.2 安装必要的依赖

除了常规的PyTorch和TensorRT外，还需要安装一些辅助工具：

bash复制pip install nvidia-pyindex
pip install tensorrt
pip install pycuda

这里最容易出问题的就是版本兼容性。我强烈建议使用官方推荐的版本组合，否则后面可能会遇到各种奇怪的错误。

3. INT8量化实现过程

3.1 校准数据集准备

INT8量化的核心是校准(calibration)过程，需要准备一个有代表性的数据集。这个数据集不需要标注，但应该覆盖实际应用中的各种场景。

我创建了一个简单的DataLoader类来处理校准图像：

python复制class CalibrationDataLoader:
    def __init__(self, data_dir, batch_size=8, img_size=640):
        self.img_files = glob.glob(f"{data_dir}/*.jpg")
        self.batch_size = batch_size
        self.img_size = img_size
        self.current_idx = 0
        
    def preprocess(self, img_path):
        img = cv2.imread(img_path)
        img = cv2.resize(img, (self.img_size, self.img_size))
        img = img.transpose(2, 0, 1).astype(np.float32)
        img /= 255.0
        return img
        
    def next_batch(self):
        if self.current_idx >= len(self.img_files):
            return None
            
        batch = []
        for _ in range(self.batch_size):
            if self.current_idx < len(self.img_files):
                img = self.preprocess(self.img_files[self.current_idx])
                batch.append(img)
                self.current_idx += 1
                
        return np.array(batch)

3.2 实现INT8校准器

TensorRT需要自定义一个校准器来实现INT8量化。这里我参考了官方示例，实现了一个简单的校准器：

python复制class YOLOv5Calibrator(trt.IInt8EntropyCalibrator2):
    def __init__(self, data_loader, cache_file=""):
        super().__init__()
        self.data_loader = data_loader
        self.cache_file = cache_file
        self.current_idx = 0
        self.device_input = cuda.mem_alloc(self.data_loader.batch_size * 3 * 640 * 640 * 4)
        
    def get_batch_size(self):
        return self.data_loader.batch_size
        
    def get_batch(self, names):
        batch = self.data_loader.next_batch()
        if batch is None:
            return None
            
        cuda.memcpy_htod(self.device_input, batch.astype(np.float32))
        return [int(self.device_input)]
        
    def read_calibration_cache(self):
        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):
        with open(self.cache_file, "wb") as f:
            f.write(cache)

4. 常见问题与解决方案

4.1 TensorRT版本兼容性问题

我最初使用的是GitHub上找到的一个开源代码，结果运行时出现各种错误。经过排查发现是TensorRT版本不兼容的问题。原代码是针对TensorRT 7.x设计的，而我用的是8.x版本。

主要修改点包括：

1. BuilderConfig的创建方式变化
2. INT8标志设置方式不同
3. 网络构建流程有调整

新版TensorRT的代码应该这样写：

python复制def build_engine(onnx_path, calib=None):
    logger = trt.Logger(trt.Logger.VERBOSE)
    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:
        if not parser.parse(f.read()):
            for error in range(parser.num_errors):
                print(parser.get_error(error))
            return None
            
    config = builder.create_builder_config()
    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)
    
    if calib is not None:
        config.set_flag(trt.BuilderFlag.INT8)
        config.int8_calibrator = calib
        
    serialized_engine = builder.build_serialized_network(network, config)
    return serialized_engine

4.2 量化后精度下降问题

INT8量化后模型精度确实会有下降，但可以通过以下方法尽量减小影响：

1. 使用更多样化的校准数据集
2. 增加校准迭代次数
3. 尝试不同的校准方法（如熵校准或最小最大校准）

在我的测试中，使用500张图片进行校准，mAP下降了约1.5%，但在推理速度上获得了2.3倍的提升。

5. 性能对比与结果分析

经过多次实验，我得到了以下对比数据：

从结果可以看出，INT8量化在几乎不影响精度的情况下，显著减小了模型体积并提升了推理速度。特别是在Jetson等边缘设备上，这种优化效果会更加明显。

6. 完整代码实现

最后，我把完整的实现代码分享出来，供大家参考：

python复制import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
import cv2
import os
import glob

class YOLOv5Calibrator(trt.IInt8EntropyCalibrator2):
    # 校准器实现，同上文
    
def build_engine(onnx_path, calib=None):
    # 引擎构建函数，同上文
    
def main():
    # 准备校准数据
    calib_data = CalibrationDataLoader("calib_images", batch_size=8)
    calib = YOLOv5Calibrator(calib_data, "calib.cache")
    
    # 构建INT8引擎
    engine = build_engine("yolov5s.onnx", calib)
    
    # 保存引擎
    with open("yolov5s_int8.engine", "wb") as f:
        f.write(engine)
        
if __name__ == "__main__":
    main()

在实际项目中，我发现INT8量化特别适合部署在资源受限的设备上。虽然过程有些曲折，但看到量化后的模型能在边缘设备上流畅运行，所有的努力都是值得的。建议大家在遇到问题时多参考TensorRT的官方示例，通常都能找到解决方案。