4、TensorRT量化工具链：PTQ流程、校准数据集准备、INT8校准器详解

说到TensorRT的量化，很多人第一反应就是“精度掉得厉害”。其实不然。我做了这么多年的部署，INT8量化只要工具链用得对，精度损失完全可以控制在1%以内。今天咱们就聊聊TensorRT的PTQ（训练后量化）这套工具链，说白了就是怎么把训练好的FP32模型，无损地压成INT8。

4.1 PTQ流程：三步走，不折腾

PTQ的流程其实不复杂。我个人习惯把它拆成三步：

1. 准备校准数据集——喂给模型一些有代表性的数据
2. 运行校准——让TensorRT统计每层的激活值分布
3. 生成INT8引擎——根据统计结果做量化

嗯，这里要注意：校准数据集的质量，直接决定了量化后的精度。我在项目中遇到过好几次，校准集选得不好，模型精度直接掉5个点以上。

4.2 校准数据集准备：别随便拿100张图就开干

校准数据集怎么准备？很多人随便从训练集里抽100张图就开干。你想想看，这样能行吗？

我给大家几个硬性标准：

• 数量：一般500~1000张就够了。太少统计不准，太多也没必要
• 多样性：要覆盖模型在真实场景中可能遇到的各种情况。比如做目标检测，光照、角度、遮挡都要有
• 代表性：校准集的数据分布要和真实部署场景一致。我曾经有个项目，校准集全是白天拍的，结果部署到夜间场景，精度直接崩了

4.3 INT8校准器（Calibrator）详解

校准器是PTQ的核心。TensorRT提供了几种校准器，我一个个说。

4.3.1 校准器类型

校准器	原理	适用场景
EntropyCalibrator	最小化量化前后的KL散度	通用场景，精度最优
MinMaxCalibrator	直接取激活值的min/max	对精度要求不高，追求速度
PercentileCalibrator	去掉一定百分比的极端值	激活值有长尾分布时

我个人最常用的是EntropyCalibrator。为什么？因为它通过KL散度找到的阈值，能最大程度保留原始分布的信息。说白了就是“砍掉尾巴，保住身子”。

4.3.2 自定义校准器实战

很多时候，TensorRT自带的校准器不够用。比如我遇到过一个模型，激活值分布特别奇怪，用EntropyCalibrator量化后精度掉了3%。这时候就需要自定义校准器。

下面是一个自定义校准器的代码示例：

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

class MyCalibrator(trt.IInt8EntropyCalibrator2):
    def __init__(self, calibration_data, batch_size=32):
        trt.IInt8EntropyCalibrator2.__init__(self)
        self.calibration_data = calibration_data
        self.batch_size = batch_size
        self.current_index = 0
        
        # 分配GPU内存
        self.device_input = cuda.mem_alloc(
            self.calibration_data[0].nbytes * self.batch_size
        )
    
    def get_batch_size(self):
        return self.batch_size
    
    def get_batch(self, names):
        if self.current_index + self.batch_size > len(self.calibration_data):
            return None
        
        batch = self.calibration_data[
            self.current_index:self.current_index + self.batch_size
        ]
        self.current_index += self.batch_size
        
        # 拷贝到GPU
        cuda.memcpy_htod(self.device_input, np.ascontiguousarray(batch))
        return [int(self.device_input)]
    
    def read_calibration_cache(self, length):
        return None
    
    def write_calibration_cache(self, cache):
        # 可以保存校准结果，下次直接加载
        with open('calibration.cache', 'wb') as f:
            f.write(cache)

4.4 校准器的工作流程

为了让大家更直观地理解，我画了一张流程图：

这张图展示了整个PTQ的流程。你看，FP32模型和校准数据集是输入，经过校准器统计分布，最后生成INT8引擎。每一步都环环相扣。

4.5 实战中的几个关键点

最后，我分享几个实战中容易踩的坑：

• 校准数据要打乱：不要按类别顺序喂数据，否则校准器会学到错误的分布
• batch_size不要太大：我一般设32或64，太大容易显存溢出
• 缓存校准结果：校准一次可能要跑几万张图，保存cache文件可以省下次校准的时间
• 验证精度：量化后一定要在验证集上跑一遍，别只看训练集精度

好了，PTQ的流程和校准器就聊到这儿。说白了，校准器就是帮你找到那个“砍多少最合适”的阈值。多试几种校准器，找到最适合你模型的那一个。

---

公众号:蓝海资料掘金营，微信deep3321