4、PTQ（训练后量化）实战：使用 TensorRT 对分类模型进行 INT8 量化

好，咱们直接进入正题。这一章，我带你手把手走一遍 TensorRT 的 INT8 量化流程。说白了，就是把一个训练好的 FP32 模型，通过 PTQ（训练后量化）变成 INT8 模型。为什么要做这个？因为 INT8 模型跑得快、占内存小，在嵌入式设备上特别香。

我个人习惯把 PTQ 分成三步：校准数据准备、量化执行、精度验证。咱们一个一个来。

4.1 校准数据集的选择与生成

INT8 量化不是凭空乱压的。它需要一小批数据来“观察”模型各层的激活值分布，然后决定怎么把 FP32 的数值映射到 INT8 的 256 个刻度上。这批数据就叫校准数据集。

校准数据集长什么样？

  • 数量不用多：一般 500~1000 张图片就够了。我见过有人用 100 张也能跑，但精度会抖得厉害。
  • 来源必须真实：一定要从训练集或验证集里抽，不能用测试集。测试集是留着最后打分的，提前用了就作弊了。
  • 分布要覆盖全：每个类别至少要有几张。比如你的分类模型有 10 类，那校准集里每类至少 50 张，别让某个类别“缺席”。
我的经验： 校准集的质量比数量重要。我曾经用 200 张精心挑选的图片，比用 1000 张随机图片效果还好。你想想看，如果随机抽的 1000 张里 800 张都是背景相似的图，那模型对“边缘案例”的量化就完全没学到。

生成校准集的代码示例：

import os
import random
from PIL import Image
import numpy as np

def generate_calibration_set(data_dir, output_file, num_samples=500):
    """
    从训练集中随机抽取图片，生成校准集文件列表
    """
    all_images = []
    for root, dirs, files in os.walk(data_dir):
        for f in files:
            if f.endswith(('.jpg', '.jpeg', '.png')):
                all_images.append(os.path.join(root, f))
    
    # 随机打乱并选取
    random.shuffle(all_images)
    selected = all_images[:num_samples]
    
    # 写入文件，每行一个路径
    with open(output_file, 'w') as f:
        for img_path in selected:
            f.write(img_path + '\n')
    
    print(f"校准集已生成，共 {len(selected)} 张图片")
    return selected

# 使用示例
generate_calibration_set('./train_data', './calib_list.txt', num_samples=500)

4.2 使用 TensorRT 进行 INT8 量化

校准集准备好了，接下来就是 TensorRT 的主场。TensorRT 的 INT8 量化需要你提供一个校准器（Calibrator），它负责读取校准图片，喂给模型，然后统计每层的激活值分布。

核心流程：

  1. 加载 FP32 模型（ONNX 或 TensorRT 格式）
  2. 创建 INT8 配置（BuilderConfig）
  3. 设置校准器（Calibrator）
  4. 构建引擎（Build Engine）
  5. 序列化保存（Serialize）

嗯，这里要注意：TensorRT 的校准器有多种策略，最常用的是 IInt8EntropyCalibrator2（熵校准）。我个人习惯用这个，因为它对分类模型比较友好。

校准器实现代码：

import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
from PIL import Image

class Calibrator(trt.IInt8EntropyCalibrator2):
    def __init__(self, calib_list, batch_size, input_shape):
        trt.IInt8EntropyCalibrator2.__init__(self)
        self.calib_list = calib_list
        self.batch_size = batch_size
        self.input_shape = input_shape  # (C, H, W)
        self.batch_idx = 0
        self.device_input = cuda.mem_alloc(batch_size * input_shape[0] * input_shape[1] * input_shape[2] * 4)
        
        # 读取所有图片路径
        with open(calib_list, 'r') as f:
            self.image_paths = [line.strip() for line in f.readlines()]
        
        self.num_batches = len(self.image_paths) // batch_size
    
    def get_batch(self, names):
        if self.batch_idx >= self.num_batches:
            return None
        
        # 读取当前 batch 的图片
        batch_images = []
        start = self.batch_idx * self.batch_size
        end = start + self.batch_size
        for i in range(start, end):
            img = Image.open(self.image_paths[i]).resize((self.input_shape[2], self.input_shape[1]))
            img = np.array(img, dtype=np.float32)
            # 归一化到 [0, 1]
            img = img / 255.0
            # 调整维度为 (C, H, W)
            img = img.transpose((2, 0, 1))
            batch_images.append(img)
        
        batch_array = np.array(batch_images, dtype=np.float32)
        cuda.memcpy_htod(self.device_input, batch_array)
        
        self.batch_idx += 1
        return [self.device_input]
    
    def get_batch_size(self):
        return self.batch_size
    
    def read_calibration_cache(self):
        return None
    
    def write_calibration_cache(self, cache):
        with open('calibration.cache', 'wb') as f:
            f.write(cache)

构建 INT8 引擎：

def build_int8_engine(onnx_path, calib_list, batch_size=32):
    logger = trt.Logger(trt.Logger.INFO)
    builder = trt.Builder(logger)
    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    parser = trt.OnnxParser(network, logger)
    
    # 解析 ONNX 模型
    with open(onnx_path, 'rb') as f:
        parser.parse(f.read())
    
    config = builder.create_builder_config()
    config.set_flag(trt.BuilderFlag.INT8)
    
    # 设置校准器
    input_shape = (3, 224, 224)  # 根据你的模型调整
    calibrator = Calibrator(calib_list, batch_size, input_shape)
    config.int8_calibrator = calibrator
    
    # 构建引擎
    engine = builder.build_engine(network, config)
    
    # 序列化保存
    with open('model_int8.engine', 'wb') as f:
        f.write(engine.serialize())
    
    print("INT8 引擎构建完成，已保存为 model_int8.engine")
    return engine
避坑指南： 我曾经在构建 INT8 引擎时卡了整整一天，最后发现是校准器的 get_batch 函数返回了错误的数据类型。TensorRT 要求返回的是 int 类型的设备内存指针，千万别传 numpy 数组进去。另外，校准图片的预处理必须和训练时完全一致，包括归一化参数、通道顺序等。

4.3 精度评估与调优

模型量化完了，不能直接上线。你得先跑一遍精度评估，看看 INT8 模型和 FP32 模型差了多少。

评估流程：

  1. 加载 FP32 模型和 INT8 模型
  2. 在验证集上分别跑推理
  3. 计算 Top-1 / Top-5 准确率
  4. 对比精度差异

精度评估代码：

import tensorrt as trt
import numpy as np
from PIL import Image
import os

def evaluate_engine(engine_path, val_dir, label_file, batch_size=32):
    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()
    
    # 准备输入输出缓冲区
    input_idx = engine.get_binding_index('input')
    output_idx = engine.get_binding_index('output')
    input_shape = engine.get_binding_shape(input_idx)
    output_shape = engine.get_binding_shape(output_idx)
    
    # 加载标签
    with open(label_file, 'r') as f:
        labels = [line.strip() for line in f.readlines()]
    
    correct = 0
    total = 0
    
    # 遍历验证集
    for class_name in os.listdir(val_dir):
        class_dir = os.path.join(val_dir, class_name)
        if not os.path.isdir(class_dir):
            continue
        
        for img_name in os.listdir(class_dir):
            img_path = os.path.join(class_dir, img_name)
            img = Image.open(img_path).resize((input_shape[3], input_shape[2]))
            img = np.array(img, dtype=np.float32) / 255.0
            img = img.transpose((2, 0, 1))
            img = np.expand_dims(img, axis=0)
            
            # 推理
            output = np.zeros(output_shape, dtype=np.float32)
            context.execute_v2([img, output])
            
            # 获取预测结果
            pred_class = np.argmax(output[0])
            true_class = int(class_name)
            
            if pred_class == true_class:
                correct += 1
            total += 1
    
    accuracy = correct / total * 100
    print(f"模型精度: {accuracy:.2f}%")
    return accuracy

精度调优策略：

问题 可能原因 解决方案
精度下降 > 1% 校准集分布不匹配 重新生成校准集，确保覆盖所有类别
某个类别精度暴跌 该类别在校准集中样本太少 为该类别增加校准样本
整体精度下降 > 3% 模型对量化敏感 尝试使用 IInt8MinMaxCalibrator 或开启 FP16 混合精度
推理结果全是同一个类别 校准器返回了空数据 检查 get_batch 函数，确保正确返回
我的调优经验： 如果 INT8 模型精度掉得厉害，别急着改代码。先检查校准集。我遇到过最坑的一次，是校准集里混入了几张损坏的图片，导致模型学到了错误的分布。后来我加了一个图片校验步骤，把损坏的图片提前过滤掉，精度直接回升了 2%。

另外，如果精度还是不满意，可以试试 逐层量化。TensorRT 允许你指定某些层保持 FP32，只量化其他层。这招对敏感层特别有效。

# 设置逐层量化
config.set_flag(trt.BuilderFlag.INT8)
# 指定某些层不量化
for layer in network:
    if layer.name in ['conv1', 'fc']:  # 这些层保持 FP32
        layer.precision = trt.float32
        layer.set_output_type(0, trt.float32)

说白了，INT8 量化就是一个“精度换速度”的买卖。你的任务就是找到那个平衡点，让精度损失在可接受范围内，同时把推理速度提上去。

好了，这一章的内容就到这。校准集怎么选、量化怎么跑、精度怎么调，你应该心里有数了。动手试试吧，遇到问题再回来翻翻这一章。

专注资料整理
← 上一章 📖 返回目录 下一章 →