3、量化工具实战：使用TensorRT进行INT8量化、校准数据集准备、量化后精度评估

好，咱们直接进入正题。

INT8量化，说白了就是把模型里那些FP32的权重和激活值，用8位整数来表示。你想想看，32位变成8位，模型体积直接缩到四分之一，推理速度也能翻倍甚至更多。在嵌入式设备上，这几乎是必走的一步。

我个人习惯用TensorRT来做这件事。它不光快，而且对INT8量化的支持非常成熟。今天我就带你走一遍完整的流程：校准数据集怎么准备、量化怎么做、精度怎么评估。

3.1 校准数据集：量化的“标尺”

很多人以为量化就是把模型扔给TensorRT，它自己就搞定了。其实没那么简单。

INT8量化需要一个关键的东西——校准数据集。它的作用是让TensorRT知道，模型在真实推理时，每一层的激活值大概分布在什么范围。有了这个分布，才能找到合适的量化阈值。

注意：校准数据集不能随便拿几张图凑数。它必须能代表真实场景的数据分布。否则量化后的精度会崩得很惨。

我在项目中遇到过一件事：有个同事拿了一堆高清风景图做校准，结果模型在监控摄像头那种低分辨率场景下，精度直接掉了5个点。后来换成监控场景的图片重新校准，问题就解决了。

那么，校准数据集怎么准备？一般遵循这几个原则：

  • 数量不用太多：500~1000张就够了。太多反而增加校准时间，收益不大。
  • 覆盖要广：各种光照、角度、目标类别都要有。别只挑“好认”的样本。
  • 预处理要一致：校准时的预处理（resize、归一化等）必须和训练时完全一致。

下面是一个典型的校准数据集准备代码：

import cv2
import os
import numpy as np

class CalibDataLoader:
    def __init__(self, image_dir, batch_size=32):
        self.image_paths = [os.path.join(image_dir, f) 
                           for f in os.listdir(image_dir) 
                           if f.endswith(('.jpg', '.png'))]
        self.batch_size = batch_size
        self.index = 0
        
    def get_batch(self):
        """返回一个batch的预处理图像"""
        batch = []
        for _ in range(self.batch_size):
            if self.index >= len(self.image_paths):
                self.index = 0  # 循环使用
            img = cv2.imread(self.image_paths[self.index])
            img = cv2.resize(img, (224, 224))
            img = img.astype(np.float32) / 255.0
            img = (img - [0.485, 0.456, 0.406]) / [0.229, 0.224, 0.225]
            batch.append(img)
            self.index += 1
        return np.array(batch)

嗯，这里要注意：get_batch 方法返回的数组形状必须是 (batch_size, C, H, W) 或者 (batch_size, H, W, C)，取决于你的模型输入格式。TensorRT默认是NCHW。

3.2 TensorRT INT8量化实战

校准数据集准备好了，接下来就是量化本身。TensorRT提供了两种方式来做INT8量化：

  1. 隐式量化：用TensorRT的API直接构建网络，设置精度为INT8。
  2. 显式量化：先用ONNX导出模型，再用TensorRT的ONNX解析器加载并量化。

我个人更推荐第二种方式。因为现在大部分模型都是用PyTorch或TensorFlow训练，转成ONNX再量化，流程更清晰，也更容易调试。

来看一个完整的量化流程：

import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit

# 创建logger和builder
logger = trt.Logger(trt.Logger.INFO)
builder = trt.Builder(logger)

# 创建网络定义
network = builder.create_network(
    1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
)

# 加载ONNX模型
parser = trt.OnnxParser(network, logger)
with open("model.onnx", "rb") as f:
    if not parser.parse(f.read()):
        for error in range(parser.num_errors):
            print(parser.get_error(error))

# 配置INT8量化
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB

# 关键：设置INT8模式并指定校准器
config.set_flag(trt.BuilderFlag.INT8)
calibrator = MyCalibrator(calib_data_loader)
config.int8_calibrator = calibrator

# 构建引擎
engine = builder.build_serialized_network(network, config)

这里有个坑，我踩过好几次——校准器必须继承自 trt.IInt8EntropyCalibrator2。TensorRT提供了几种校准算法：

校准器类型 特点 适用场景
EntropyCalibrator 基于信息熵，精度较高 通用场景，推荐首选
EntropyCalibrator2 改进版，更稳定 大部分模型都适用
MinMaxCalibrator 基于最大值最小值，速度快 对精度要求不高的场景
LegacyCalibrator 旧版，不推荐 兼容老模型

我个人习惯用 EntropyCalibrator2。它在精度和速度之间平衡得最好。

下面是一个完整的校准器实现：

import tensorrt as trt
import pycuda.driver as cuda
import numpy as np

class MyCalibrator(trt.IInt8EntropyCalibrator2):
    def __init__(self, data_loader, batch_size=32):
        super().__init__()
        self.data_loader = data_loader
        self.batch_size = batch_size
        self.device_input = cuda.mem_alloc(
            batch_size * 3 * 224 * 224 * 4  # float32
        )
        
    def get_batch_size(self):
        return self.batch_size
    
    def get_batch(self, names):
        """返回一个batch的数据到GPU"""
        batch = self.data_loader.get_batch()
        if batch is None:
            return None
        # 拷贝到GPU
        cuda.memcpy_htod(self.device_input, batch.astype(np.float32))
        return [int(self.device_input)]
    
    def read_calibration_cache(self):
        """如果有缓存，直接读取，避免重复校准"""
        return None
    
    def write_calibration_cache(self, cache):
        """保存校准缓存，下次可以直接用"""
        with open("calibration.cache", "wb") as f:
            f.write(cache)
小技巧：校准缓存文件可以保存下来。下次构建引擎时，如果模型没变，直接读取缓存，能省掉校准的时间。我在项目中经常这么干，尤其是反复调试的时候。

3.3 量化后精度评估

量化完了，心里最没底的就是精度。到底掉了多少？能不能接受？

评估精度其实很简单：用量化后的模型跑一遍验证集，跟FP32模型的精度对比一下。

但这里有个细节——不能只对比一个指标。比如分类任务，除了Top-1准确率，还要看Top-5。检测任务要看mAP。我见过一个项目，Top-1只掉了0.5%，但某个关键类别的召回率掉了8%，上线后直接翻车。

下面是一个精度对比的代码框架：

import numpy as np
from tqdm import tqdm

def evaluate_model(engine_path, val_loader, num_classes=1000):
    """评估量化后模型的精度"""
    # 加载引擎
    with open(engine_path, "rb") as f:
        engine_data = f.read()
    
    runtime = trt.Runtime(logger)
    engine = runtime.deserialize_c_engine(engine_data)
    context = engine.create_execution_context()
    
    # 分配输入输出缓冲区
    input_idx = 0
    output_idx = 1
    input_shape = (1, 3, 224, 224)
    output_shape = (1, num_classes)
    
    # 评估
    correct = 0
    total = 0
    
    for images, labels in tqdm(val_loader):
        # 推理
        outputs = inference(context, engine, images)
        preds = np.argmax(outputs, axis=1)
        correct += (preds == labels.numpy()).sum()
        total += len(labels)
    
    accuracy = correct / total
    return accuracy

# 对比FP32和INT8
fp32_acc = evaluate_model("model_fp32.engine", val_loader)
int8_acc = evaluate_model("model_int8.engine", val_loader)

print(f"FP32精度: {fp32_acc:.4f}")
print(f"INT8精度: {int8_acc:.4f}")
print(f"精度损失: {(fp32_acc - int8_acc)*100:.2f}%")

一般来说，INT8量化后精度损失在1%以内是可以接受的。如果超过2%，就要排查问题了。

精度下降的常见原因：
  • 校准数据集分布偏差太大
  • 模型中有对量化敏感的层（如某些激活函数）
  • 校准图片数量太少（少于200张）
  • 预处理不一致（校准和推理用的预处理不同）

我曾经遇到过一个案例：模型量化后精度掉了3%。排查了半天，发现是校准数据集的归一化参数写错了——用了mean=[0.5,0.5,0.5]，但训练时用的是mean=[0.485,0.456,0.406]。改过来之后，精度损失降到了0.3%。

所以，校准数据集的预处理必须和训练时完全一致。这个坑，我建议你提前记下来。

3.4 总结一下

INT8量化不是玄学，它有一套清晰的流程：

  1. 准备一个有代表性的校准数据集（500~1000张）
  2. 实现校准器，继承 IInt8EntropyCalibrator2
  3. 用TensorRT构建INT8引擎
  4. 在验证集上对比FP32和INT8的精度

如果精度损失在可接受范围内，恭喜你，模型已经成功瘦身加速了。如果损失过大，别慌，先检查校准数据集和预处理，大概率问题出在那里。

下一章，我们会聊到更进阶的话题——量化敏感层分析与混合精度量化。到时候我会分享一些“抢救”精度损失的实战技巧。

← 上一章 📖 返回目录 下一章 →