TensorRT 推理优化：FP16、INT8 与层融合

TensorRT 推理优化：FP16、INT8 与层融合

TensorRT 是 NVIDIA 提供的高性能深度学习推理优化器和运行时引擎。它能够将训练好的模型转换为极致优化的推理引擎，在保持精度的同时大幅降低延迟、提高吞吐量。本教程面向初学者，聚焦三大核心优化技术：FP16 精度、INT8 量化以及层融合，手把手带你理解原理并付诸实践。

1. 为什么要进行推理优化？

训练出来的模型通常使用 FP32（单精度浮点）进行参数存储和计算，这在推理阶段会导致：

• 高显存占用：模型权重和中间激活占满 GPU 内存。
• 高延迟：大量的计算操作无法充分利用 GPU 硬件。
• 低吞吐：在线服务无法满足实时性要求。

TensorRT 通过对模型图进行重构和精度降低，可以在几乎不损失精度的情况下，达到数倍的性能提升。

2. 环境准备与基本流程

在开始之前，请确保已安装合适版本的 TensorRT 和 CUDA。可通过 NVIDIA 官网下载或使用 pip 安装：

pip install nvidia-tensorrt

典型的工作流程分为三个阶段：

1. 模型转换：将 PyTorch / TensorFlow / ONNX 模型转换为 TensorRT 引擎。
2. 构建引擎：选择优化配置（精度、层融合策略）并生成 .trt 引擎文件。
3. 执行推理：加载引擎并运行高速推理。

我们将围绕 ONNX 路径进行演示，这是兼容性最好的方式。

3. 精度优化：FP16 推理

FP16（半精度浮点）使用 16 位存储一个数，相比 FP32 能带来接近 2 倍 的带宽和计算吞吐提升，并且显存占用减半。现代 GPU（如 Volta、Turing、Ampere 架构）的 Tensor Core 对 FP16 有专门硬件加速。

3.1 开启 FP16 模式

在 TensorRT 中，构建引擎时只需设置一个标志：

import tensorrt as trt

def build_engine_fp16(onnx_file_path, engine_file_path):
    logger = trt.Logger(trt.Logger.WARNING)
    builder = trt.Builder(logger)
    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    parser = trt.OnnxParser(network, logger)
    
    with open(onnx_file_path, 'rb') as f:
        parser.parse(f.read())
    
    config = builder.create_builder_config()
    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB workspace
    config.set_flag(trt.BuilderFlag.FP16)   # 启用 FP16
    
    engine = builder.build_engine(network, config)
    with open(engine_file_path, 'wb') as f:
        f.write(engine.serialize())

关键点：

• 只需 config.set_flag(trt.BuilderFlag.FP16) 一行即可启用。
• TensorRT 会自动选择哪些层用 FP16 计算，对精度敏感的层仍保留 FP32，这个过程称为 自动混合精度。
• 支持 FP16 的 GPU 架构须在 Compute Capability 7.0 以上。

3.2 验证精度与性能

构建完成后，使用简单的推理循环对比 FP32 和 FP16 的输出差异以及延迟。通常 FP16 的输出与 FP32 的相对误差小于 1e-3，对大多数视觉和 NLP 模型足够。

# 示例：性能测试
import time
import numpy as np

def benchmark(engine):
    # 分配输入输出缓冲区，执行多次推理取平均
    ...

4. 精度优化：INT8 量化

INT8 量化将权重和激活值从 32 位浮点映射到 8 位整数，理论性能可达 FP32 的 4 倍，显存占用降至 1/4。但直接量化会引入显著误差，因此需要校准来确定动态范围。

4.1 校准原理

TensorRT 的 INT8 校准分为两种：

• 熵校准：最小化原始分布和量化分布的 KL 散度。
• Legacy（最小值/最大值）：直接使用绝对值最大值作为缩放因子。

校准需要少量有代表性的输入数据（通常 500-1000 张图片），TensorRT 会统计各层的激活值范围。

4.2 实现 INT8 引擎

我们需要实现 IInt8Calibrator 接口或使用便捷的 trt.IInt8EntropyCalibrator2。下面的例子使用图像批次进行熵校准：

class MyInt8Calibrator(trt.IInt8EntropyCalibrator2):
    def __init__(self, data_loader, cache_file='calibration.cache'):
        trt.IInt8EntropyCalibrator2.__init__.__init__(self)
        self.data_loader = data_loader
        self.cache_file = cache_file
        self.batch_iterator = iter(data_loader)
        # 根据实际输入定义设备缓冲区
        self.device_input = None

    def get_batch_size(self):
        return 1  # 每次返回一个批次

    def get_batch(self, names):
        try:
            batch = next(self.batch_iterator)
            # 将数据复制到 GPU
            if not self.device_input:
                self.device_input = cuda.mem_alloc(batch.nbytes)
            cuda.memcpy_htod(self.device_input, batch.ravel())
            return [int(self.device_input)]
        except StopIteration:
            return None

    def read_calibration_cache(self):
        if os.path.exists(self.cache_file):
            with open(self.cache_file, 'rb') as f:
                return f.read()

    def write_calibration_cache(self, cache):
        with open(self.cache_file, 'wb') as f:
            f.write(cache)

构建引擎时，配置 INT8 并传入校准器：

config.set_flag(trt.BuilderFlag.INT8)
calibrator = MyInt8Calibrator(data_loader)
config.int8_calibrator = calibrator

当校准缓存存在时，后续构建会跳过校准过程，直接使用缓存。

4.3 动态范围与精度调优

• 对于某些敏感层，TensorRT 可能会回退到 FP16 或 FP32，这是自动的。
• 如果 INT8 精度损失较大，可以尝试：增大校准数据集、使用 IInt8LegacyCalibrator 或启用 OBEY_PRECISION_CONSTRAINTS 强制完全 INT8（慎用）。

5. 图优化：层融合

层融合是 TensorRT 最重要的图优化技术之一，它将多个连续的小操作合并成一个优化的大操作，从而减少内存带宽消耗和 kernel 启动开销。

5.1 常见融合模式

TensorRT 会在构建引擎时自动执行这些融合，无需用户干预。但理解其背后的逻辑有助于模型设计。

5.2 如何确认融合生效？

可以通过设置日志级别查看优化细节：

logger = trt.Logger(trt.Logger.VERBOSE)

输出中会显示哪些层被融合（例如：Fusing conv1 + bn1 + relu1 into CBR）。还可以使用 trtexec 工具的 --dumpProfile 来导出层时间并分析融合情况。

5.3 手动控制融合

高级用户可通过 Layer Precision 或 Tactic Source 微调，但初学者一般无需操作。只需注意：批归一化（BN）和卷积的融合需要合并 BN 参数到卷积权重中，这要求模型在导出 ONNX 前已将 BN 层折叠（如 PyTorch 中使用 torch.onnx.export 时设置 training=torch.onnx.TrainingMode.EVAL 或手动执行 fuse_conv_bn）。

6. 综合优化策略：FP16 + INT8 + 层融合

实际部署时，常常同时开启多种优化以获得最大收益。推荐配置顺序：

1. 首先使用 FP16，几乎无精度损失，收益明显。
2. 若需更高吞吐，尝试 INT8，并进行校准验证精度。
3. 确保模型结构利于融合，如使用 Conv2d --> BatchNorm2d --> ReLU 的标准块，避免在关键路径插入过多的自定义操作。

以下为构建支持双重精度和自动融合的通用代码骨架：

def build_engine(onnx_path, engine_path, use_fp16=True, use_int8=False, calibrator=None):
    builder = trt.Builder(trt.Logger(trt.Logger.WARNING))
    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    parser = trt.OnnxParser(network, logger)
    
    with open(onnx_path, 'rb') as f:
        parser.parse(f.read())

    config = builder.create_builder_config()
    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)

    if use_int8 and calibrator:
        config.set_flag(trt.BuilderFlag.INT8)
        config.int8_calibrator = calibrator
    elif use_fp16:
        config.set_flag(trt.BuilderFlag.FP16)

    engine = builder.build_engine(network, config)
    with open(engine_path, 'wb') as f:
        f.write(engine.serialize())

7. 常见问题与解决

Q1: ONNX 模型导入时某些操作不支持？

A: TensorRT 对 ONNX 算子有较完整支持，但某些新算子可能缺失。可尝试升级 TensorRT 版本，或使用 Plugin 注册自定义算子，也可以将模型拆分为子图，用原生 CUDA 补全。

Q2: INT8 校准过程中显存不足？

A: 减小校准批次大小，或在校准时限制使用的图像数量。

Q3: FP16 或 INT8 模型输出完全错误？

A: 检查 GPU 驱动和 CUDA 版本是否匹配，确保 TensorRT 版本支持你的 GPU 架构。还可尝试关闭 CPU 多线程（trt.BuilderFlag.DISABLE_TIMING_CACHE）来排除缓存问题。

Q4: 层融合未带来预期提速？

A: 可能是因为模型已经高度优化，或者计算瓶颈不在融合的层上（如注意力机制中的大量矩阵乘法）。使用 NVIDIA Nsight Systems 定位真正的瓶颈。

8. 总结

TensorRT 将 精度降低（FP16/INT8）与 计算图优化（层融合）相结合，为模型部署提供了极致的性能加速。掌握这三项技术后，你将能轻松将模型推理成本降低 50% - 75%，同时保持可用精度。接下来，你可以尝试用自己的模型导出 ONNX，逐步开启 FP16 和 INT8，感受推理速度的飞跃。

继续探索：参考 TensorRT 官方文档中的 trtexec 命令行工具，它无需编写代码即可快速转换和评测模型。