📚 课程目录第4章 TensorRT量化工具

4、TensorRT量化工具：从FP16到INT8的实战之路

说到模型部署，TensorRT几乎是绕不开的一个工具。我个人习惯把它叫做「英伟达家的性能榨汁机」——你训练好的模型，经过它一处理，推理速度能翻好几倍。

这一章，我们就来聊聊TensorRT的量化能力。说白了，就是怎么让YOLOv11跑得更快，同时尽量不掉精度。

4.1 TensorRT是什么？

TensorRT是NVIDIA推出的深度学习推理优化引擎。它不是用来训练的，专门负责把训练好的模型「翻译」成能在GPU上高效运行的版本。

我刚开始接触时也犯过迷糊：明明PyTorch模型能跑，为什么还要多此一举？后来在项目里遇到一个实时检测需求，原模型跑不到30帧，用TensorRT优化后直接飙到90帧。嗯，从那以后我就再也没跳过这一步。

TensorRT的核心能力包括：

层融合：把多个小算子合并成一个，减少Kernel启动开销
精度校准：支持FP16和INT8量化，用更少的位宽换速度
内存优化：复用显存，减少碎片
自动调优：为你的GPU型号选择最快的Kernel

一句话总结：TensorRT = 模型加速器 + 量化工具 + 部署引擎。你给它一个模型，它还你一个飞快的引擎。

4.2 FP16量化与INT8量化

量化，说白了就是用更少的比特数来表示模型参数和激活值。FP32是32位浮点，FP16是16位，INT8是8位整型。位宽越小，计算越快，显存占用也越小。

但代价是什么？精度损失。你想想看，用更少的数字去近似原来的值，肯定会有误差。

4.2.1 FP16量化

FP16量化是最温和的。它把FP32的权重和激活值直接转成FP16。我做过测试，YOLOv11s模型在FP16下，mAP下降通常不到0.5%，但推理速度能提升1.5到2倍。

为什么损失这么小？因为FP16的动态范围虽然比FP32窄，但对于大多数网络层的数值分布来说，已经够用了。

我的经验：FP16量化几乎可以无脑用。除非你的模型对精度极度敏感（比如医疗影像），否则FP16是性价比最高的选择。

4.2.2 INT8量化

INT8量化就激进多了。它把32位浮点数压缩到8位整型，数值范围从约10^38缩小到-128到127。这中间的信息损失，你想想看有多大。

但好处也明显：

推理速度提升3-4倍
显存占用减少75%
带宽需求大幅降低

不过，INT8量化不能直接硬转。你需要一个「校准」过程，来找到最佳的量化参数。

量化类型	位宽	速度提升	精度损失	适用场景
FP32	32位	1x（基准）	无	训练、精度优先
FP16	16位	1.5-2x	<0.5%	大多数部署场景
INT8	8位	3-4x	1-3%	实时性要求极高

4.3 使用Calibration数据集

INT8量化需要一个关键步骤——校准（Calibration）。

为什么要校准？因为INT8的数值范围有限，你得知道模型在真实数据上的激活值分布，才能决定怎么把FP32映射到INT8上。

校准数据集，就是用来做这件事的。它不需要很大，几百到一千张图片就够了。但有一个要求：数据分布要和你的实际部署场景一致。

我曾经在一个项目中吃过亏：用COCO的图片做校准，结果部署到工厂产线上，检测精度直接掉了5个点。为什么？因为产线图片的光照、背景和COCO差太远了。校准数据不匹配，量化参数就是错的。

4.3.1 校准流程

准备校准数据：从你的验证集或训练集中，随机抽取500-1000张图片
前向推理：用FP32模型跑一遍校准数据，记录每一层的激活值分布
计算量化参数：根据分布，找到每个Tensor的缩放因子（scale）和零点（zero point）
应用量化：用这些参数把模型转成INT8

避坑指南：校准数据不要用训练集里精度最高的那批图片。我见过有人专门挑「好图」做校准，结果模型在普通图片上表现很差。校准数据要「平均」，不要「偏科」。

4.3.2 校准算法

TensorRT提供了几种校准算法：

Entropy Calibration：基于KL散度，最小化量化前后的信息损失。最常用，效果也最好
MinMax Calibration：直接取激活值的最大值和最小值作为范围。简单粗暴，但容易受异常值影响
Percentile Calibration：去掉极端值，用百分位数确定范围。适合有少量异常值的场景

我个人习惯先用Entropy，如果精度下降太多，再换Percentile试试。

4.4 TensorRT模型构建与序列化

好了，理论讲完了，我们来点实际的。怎么把YOLOv11模型构建成TensorRT引擎，并保存成文件？

4.4.1 构建流程

TensorRT的构建过程，说白了就是三步：

创建Builder：初始化TensorRT的构建器
定义网络：加载ONNX模型，设置输入输出
配置参数：设置精度、工作空间、校准器等
构建引擎：调用buildSerializedNetwork生成引擎

下面是一个完整的代码示例：

import tensorrt as trt

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

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

# 加载ONNX模型
parser = trt.OnnxParser(network, logger)
with open("yolov11s.onnx", "rb") as f:
    parser.parse(f.read())

# 配置构建参数
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB

# 设置FP16量化
if builder.platform_has_fast_fp16:
    config.set_flag(trt.BuilderFlag.FP16)
    print("FP16量化已启用")

# 设置INT8量化（需要校准器）
if builder.platform_has_fast_int8:
    config.set_flag(trt.BuilderFlag.INT8)
    # 这里需要设置校准器，后面会讲
    print("INT8量化已启用")

# 构建序列化引擎
serialized_engine = builder.build_serialized_network(network, config)

# 保存到文件
with open("yolov11s_fp16.engine", "wb") as f:
    f.write(serialized_engine)

print("TensorRT引擎构建完成！")

4.4.2 INT8校准器的实现

如果你要用INT8，需要自己实现一个校准器。TensorRT提供了一个基类，你只需要继承它，实现读取数据的方法。

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

class YOLOCalibrator(trt.IInt8EntropyCalibrator2):
    def __init__(self, calib_images, batch_size=32):
        trt.IInt8EntropyCalibrator2.__init__(self)
        self.calib_images = calib_images
        self.batch_size = batch_size
        self.current_index = 0
        # 分配GPU内存
        self.device_input = cuda.mem_alloc(
            batch_size * 3 * 640 * 640 * 4
        )

    def get_batch_size(self):
        return self.batch_size

    def get_batch(self, names):
        if self.current_index + self.batch_size > len(self.calib_images):
            return None
        
        # 读取并预处理一批图片
        batch = []
        for i in range(self.batch_size):
            img = Image.open(
                self.calib_images[self.current_index + i]
            ).resize((640, 640))
            img = np.array(img, dtype=np.float32) / 255.0
            img = img.transpose(2, 0, 1)  # HWC -> CHW
            batch.append(img)
        
        batch = np.stack(batch, axis=0)
        self.current_index += self.batch_size
        
        # 拷贝到GPU
        cuda.memcpy_htod(self.device_input, batch)
        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)

小技巧：校准器跑完后会生成一个cache文件。下次构建时可以直接读取，不用重新跑校准。我一般会把cache文件保存下来，方便团队其他人复用。

4.4.3 序列化与反序列化

构建好的引擎是一个序列化的二进制文件（.engine）。部署时，你只需要加载它，不需要再跑一遍构建流程。

# 反序列化引擎
with open("yolov11s_fp16.engine", "rb") as f:
    engine_data = f.read()

runtime = trt.Runtime(logger)
engine = runtime.deserialize_cuda_engine(engine_data)

# 创建执行上下文
context = engine.create_execution_context()

# 现在就可以用context进行推理了
print("引擎加载成功，准备推理...")

这里要注意：引擎是和GPU型号绑定的。你在A100上构建的引擎，不能直接拿到T4上用。我刚开始做部署时踩过这个坑，在开发机上构建好引擎，拿到客户现场发现跑不了。嗯，从那以后我都是现场构建，或者提前确认好GPU型号。

4.5 本章小结

TensorRT的量化工具链，说白了就是一套「用精度换速度」的工程方法。FP16适合大多数场景，INT8适合对实时性有极致要求的项目。校准数据集是INT8量化的关键，选对了数据，精度损失可以控制在1%以内。

构建和序列化引擎的流程，我建议你写成一个脚本，每次更新模型后自动跑一遍。手动操作太容易出错了。