4. TensorRT INT8量化

任务目标

本案例使用TensorRT对PyTorch预训练的ResNet-18模型执行INT8量化，生成INT8 TensorRT引擎，并对比FP32与INT8引擎的推理延迟与精度差异，验证INT8量化在推理加速上的实际收益。

环境依赖

• TensorRT 8.x+（含trtexec或Python API）
• PyTorch 1.10+（仅用于导出ONNX）
• onnx 1.12+、pycuda（用于GPU内存管理）
• GPU：NVIDIA Turing架构及以上（支持INT8）

实现原理与步骤

INT8量化需提供校准数据集以统计各层激活值的动态范围。本案例采用TensorRT的熵校准（Entropy Calibration）方式，从验证集中随机选取500张图片作为校准数据。整体流程为：PyTorch模型→ONNX→TensorRT FP32引擎→TensorRT INT8引擎→性能与精度对比。

关键完整代码

import torch
import torchvision.models as models
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
from PIL import Image
from torchvision import transforms

# ---------- 1. PyTorch → ONNX ----------
def export_onnx():
    model = models.resnet18(pretrained=True).eval().cuda()
    dummy = torch.randn(1, 3, 224, 224).cuda()
    torch.onnx.export(model, dummy, "resnet18.onnx",
                      input_names=["input"], output_names=["output"],
                      dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}})
    print("ONNX导出完成")

# ---------- 2. 校准数据加载器 ----------
class CalibDataLoader:
    def __init__(self, num_calib=500):
        self.data = []
        tf = transforms.Compose([
            transforms.Resize(256), transforms.CenterCrop(224),
            transforms.ToTensor(), transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])
        ])
        # 使用随机噪声模拟校准数据（实际应用请替换为真实验证集图片）
        for _ in range(num_calib):
            img = torch.randn(3, 224, 224).numpy().astype(np.float32)
            self.data.append(img)
        self.idx = 0

    def __call__(self):
        batch = np.stack(self.data[self.idx:self.idx+50])  # 每次50张
        self.idx += 50
        return [batch]

# ---------- 3. TensorRT引擎构建 ----------
def build_engine(onnx_path, mode="FP32", calib_loader=None):
    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_path, "rb") as f:
        assert parser.parse(f.read()), "ONNX解析失败"

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

    if mode == "INT8":
        assert builder.platform_has_fast_int8, "GPU不支持INT8"
        config.set_flag(trt.BuilderFlag.INT8)
        config.int8_calibrator = calib_loader  # 自定义校准器需继承trt.IInt8Calibrator
    else:
        config.set_flag(trt.BuilderFlag.FP16)  # FP32引擎也启用FP16加速对比

    # 简化：使用trtexec方式构建（推荐）
    # 实际代码中建议调用trtexec命令行，此处仅展示API流程
    serialized = builder.build_serialized_network(network, config)
    return serialized

# ---------- 4. 推理与计时 ----------
def infer(engine_path, input_data):
    runtime = trt.Runtime(trt.Logger(trt.Logger.WARNING))
    with open(engine_path, "rb") as f:
        engine = runtime.deserialize_cuda_engine(f.read())
    context = engine.create_execution_context()
    # 分配GPU内存并执行推理（省略详细绑定代码）
    # 返回推理结果和耗时
    return np.random.rand(1000)  # 示意返回

if __name__ == "__main__":
    export_onnx()
    # 实际构建请使用trtexec命令，见下方说明

运行与验证

推荐使用trtexec命令行工具完成引擎构建与性能测试，避免手写复杂校准器代码：

# 构建FP32引擎
trtexec --onnx=resnet18.onnx --saveEngine=resnet_fp32.engine --fp16

# 构建INT8引擎（使用校准数据）
trtexec --onnx=resnet18.onnx --saveEngine=resnet_int8.engine --int8 \
        --calib=/path/to/calib_images/ --calibBatchSize=50

# 性能测试
trtexec --loadEngine=resnet_fp32.engine --shapes=input:1x3x224x224
trtexec --loadEngine=resnet_int8.engine --shapes=input:1x3x224x224

验证判定标准：

1. INT8引擎成功构建，无报错；
2. trtexec输出的INT8引擎推理延迟（latency）显著低于FP32引擎（通常提升1.5～3倍）；
3. 使用相同输入图片分别运行两个引擎，输出张量的Top-1类别一致（精度损失<1%）。

若INT8引擎构建失败，请检查GPU架构（需Turing+）并确认校准图片路径有效。实际部署时，请将校准数据替换为真实验证集图片，而非代码中的随机噪声。