4、量化工具实战：使用TensorRT进行INT8量化、使用ONNX Runtime进行量化、量化后的模型验证

好，咱们进入实战环节了。前面讲了那么多量化的理论，什么对称非对称、校准数据集、KL散度……说实话，光看理论容易晕。我自己的经验是，先动手跑通一个流程，再回头理解理论，会轻松很多。

这一章，咱们就真刀真枪地干。我会带你走两条主流路线：TensorRT的INT8量化和ONNX Runtime的量化。最后，咱们还得验证一下量化后的模型到底行不行。

4.1 用TensorRT做INT8量化

TensorRT做INT8量化，核心就三步：准备校准数据、构建校准器、生成引擎。听起来简单，但坑不少。

先看代码。假设你有一个训练好的FP32模型，比如ResNet-50。

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

# 1. 创建builder和config
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(TRT_LOGGER)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
config = builder.create_builder_config()

# 2. 设置INT8模式
config.set_flag(trt.BuilderFlag.INT8)

# 3. 设置校准器 —— 这是关键
class MyCalibrator(trt.IInt8EntropyCalibrator2):
    def __init__(self, calibration_files, batch_size, input_shape):
        trt.IInt8EntropyCalibrator2.__init__(self)
        self.batch_size = batch_size
        self.input_shape = input_shape
        self.calibration_files = calibration_files
        self.batch_idx = 0
        self.device_input = cuda.mem_alloc(trt.volume(input_shape) * trt.float32.itemsize)

    def get_batch_size(self):
        return self.batch_size

    def get_batch(self, names):
        if self.batch_idx * self.batch_size >= len(self.calibration_files):
            return None
        # 加载并预处理图像
        batch = self._load_batch()
        cuda.memcpy_htod(self.device_input, batch.astype(np.float32).ravel())
        self.batch_idx += 1
        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. 设置校准数据集路径
calib_files = ['img1.jpg', 'img2.jpg', ...]  # 至少500张
calibrator = MyCalibrator(calib_files, batch_size=8, input_shape=(3, 224, 224))
config.int8_calibrator = calibrator

# 5. 构建引擎
serialized_engine = builder.build_serialized_network(network, config)
with open('resnet50_int8.engine', 'wb') as f:
    f.write(serialized_engine)

核心要点：校准数据集必须来自真实推理场景。我见过有人拿ImageNet的验证集做校准，结果部署到自己的业务数据上，精度掉了5个点。说白了，校准数据要跟你的实际输入分布一致。

我的经验：校准图片数量一般500-1000张就够了。太少，KL散度算不准；太多，校准时间太长，收益却不大。

4.2 用ONNX Runtime做量化

ONNX Runtime的量化方式更灵活。它支持两种模式：动态量化和静态量化。我个人更推荐静态量化，因为推理速度更快。

来看代码。假设你有一个ONNX模型。

import onnx
from onnxruntime.quantization import quantize_static, QuantType
from onnxruntime.quantization import CalibrationMethod

# 1. 定义校准数据读取器
class CalibrationDataReader:
    def __init__(self, image_paths, batch_size=1):
        self.image_paths = image_paths
        self.batch_size = batch_size
        self.iter = 0

    def get_next(self):
        if self.iter * self.batch_size >= len(self.image_paths):
            return None
        batch = self._preprocess_batch()
        self.iter += 1
        return {'input': batch}

    def _preprocess_batch(self):
        # 这里做你的预处理：resize, normalize等
        pass

# 2. 执行静态量化
calib_reader = CalibrationDataReader(['img1.jpg', 'img2.jpg', ...])
quantized_model = quantize_static(
    model_input='resnet50.onnx',
    model_output='resnet50_int8.onnx',
    calibration_data_reader=calib_reader,
    quant_format=QuantType.QInt8,
    per_channel=True,
    calibration_method=CalibrationMethod.MinMax  # 或 Entropy
)

print("量化完成！")

注意：ONNX Runtime的量化默认使用MinMax校准方法。如果你的模型对异常值敏感，建议改用Entropy方法。我曾经在一个目标检测模型上踩过坑，MinMax导致边界框回归精度下降，换成Entropy就好了。

4.3 量化后的模型验证

模型量化完了，不能直接上线。你得验证三件事：精度、速度、稳定性。

4.3.1 精度验证

这个最简单，跑一遍测试集，对比FP32和INT8的指标。

def evaluate_model(model_path, test_loader, is_int8=False):
    # 加载模型
    if is_int8:
        # TensorRT或ONNX Runtime加载INT8模型
        pass
    else:
        # 加载FP32模型
        pass

    total = 0
    correct = 0
    for images, labels in test_loader:
        outputs = model(images)
        _, predicted = torch.max(outputs, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

    accuracy = 100 * correct / total
    return accuracy

fp32_acc = evaluate_model('resnet50_fp32.onnx', test_loader, is_int8=False)
int8_acc = evaluate_model('resnet50_int8.onnx', test_loader, is_int8=True)

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

行业经验：一般来说，INT8量化后精度损失在0.5%-1%以内是正常的。如果超过2%，你得检查校准数据是否合理，或者考虑使用混合精度（部分层保持FP16）。

4.3.2 速度验证

速度提升是量化的核心目标。我习惯用1000次推理取平均耗时。

import time

def benchmark(model, input_tensor, iterations=1000):
    # 预热
    for _ in range(100):
        _ = model(input_tensor)

    # 正式测试
    start = time.time()
    for _ in range(iterations):
        _ = model(input_tensor)
    end = time.time()

    avg_latency = (end - start) / iterations * 1000  # 毫秒
    return avg_latency

fp32_latency = benchmark(fp32_model, test_input)
int8_latency = benchmark(int8_model, test_input)

print(f"FP32延迟: {fp32_latency:.2f} ms")
print(f"INT8延迟: {int8_latency:.2f} ms")
print(f"加速比: {fp32_latency / int8_latency:.2f}x")

我的经验：在GPU上，INT8量化通常能带来2-4倍的加速。但在CPU上，效果可能没那么明显，因为CPU的INT8指令集支持不如GPU完善。如果你在CPU上部署，建议先查一下硬件是否支持VNNI指令。

4.3.3 稳定性验证

这个容易被忽略。量化后的模型，对同一张图片多次推理，结果应该完全一致。

def check_determinism(model, input_tensor, runs=10):
    outputs = []
    for _ in range(runs):
        out = model(input_tensor)
        outputs.append(out)

    # 检查所有输出是否一致
    for i in range(1, runs):
        if not np.allclose(outputs[0], outputs[i], atol=1e-5):
            return False
    return True

is_stable = check_determinism(int8_model, test_input)
print(f"模型稳定性: {'稳定' if is_stable else '不稳定，请检查！'}")

曾经踩过的坑：有一次我量化后的模型，精度和速度都达标了，但上线后偶尔会出现异常输出。排查了两天，发现是校准数据读取器没有固定随机种子，导致每次校准结果不一样。嗯，从此以后，我所有跟随机相关的操作都固定seed。

4.4 量化工具对比总结

对比维度 TensorRT INT8 ONNX Runtime INT8
硬件支持 仅NVIDIA GPU CPU + 部分GPU
校准方法 Entropy、MinMax、Percentile MinMax、Entropy
量化粒度 逐张量、逐通道 逐张量、逐通道
推理速度 极快（GPU优化） 较快（跨平台）
易用性 中等（需写校准器） 简单（API封装好）
适用场景 高性能GPU推理 跨平台部署

选哪个？我的建议是：如果你在NVIDIA GPU上部署，优先用TensorRT，性能最优。如果你需要跨平台（比如Windows、Linux、ARM），或者模型来自不同框架，ONNX Runtime更省心。

好了，这一章的内容就到这。量化工具实战，说白了就是三步：准备校准数据、执行量化、验证结果。每一步都有坑，但只要你按照我上面说的流程走，基本不会出大问题。下一章，咱们聊聊量化后的模型如何部署到嵌入式设备上，那才是真正的硬仗。

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