4、TensorRT引擎构建：从ONNX/PyTorch模型导出、精度校准（FP16/INT8）、动态形状（Dynamic Shape）配置

TensorRT引擎构建是将训练好的深度学习模型转化为可高效推理的部署形态的核心环节。本章聚焦于从PyTorch/ONNX模型到TensorRT引擎的完整构建链路，重点解决三个关键问题：如何正确导出模型、如何选择并实施精度校准策略、以及如何配置动态形状以应对多模型共享GPU资源时的输入尺寸变化。这三者直接决定了最终引擎的推理速度、显存占用和灵活性，是多模型调度场景下必须掌握的基础能力。

4.1 模型导出：从PyTorch到ONNX的工程化路径

模型导出是TensorRT引擎构建的第一步，其质量直接决定了后续所有环节的可行性。PyTorch模型不能直接被TensorRT消费，必须通过ONNX（Open Neural Network Exchange）作为中间表示。这一过程看似简单，但在实际工程中，算子映射失败、动态轴丢失、常量折叠错误等问题层出不穷。核心原理在于：PyTorch的torch.onnx.export函数会追踪模型的前向计算图，将其转换为ONNX的算子序列，但这一转换并非无损——某些PyTorch原生算子（如torch.einsum、torch.topk）在ONNX中可能没有直接对应，需要分解或替换。

操作流程上，推荐使用以下步骤：首先，固定模型的推理模式（model.eval()）并关闭梯度计算；其次，准备一个具有代表性的输入张量（形状、数据类型需与真实推理一致）；然后，调用torch.onnx.export并设置opset_version（建议≥13），同时通过dynamic_axes参数标记动态维度；最后，使用onnx.checker和onnxruntime进行验证，确保ONNX模型的输出与PyTorch原始输出在数值上一致（误差通常应小于1e-4）。

一个典型的导出代码示例如下：

import torch
import torch.onnx

# 假设model是训练好的PyTorch模型
model.eval()
dummy_input = torch.randn(1, 3, 224, 224, device='cuda')

torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    opset_version=17,
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={
        "input": {0: "batch_size", 2: "height", 3: "width"},
        "output": {0: "batch_size"}
    }
)
print("ONNX导出完成")

在实战中，我们曾遇到一个典型的失败案例：某目标检测模型包含torchvision.ops.nms算子，该算子在ONNX中无直接映射。解决方案是将其替换为自定义ONNX算子，并在TensorRT中通过插件（Plugin）实现。这一过程需要额外编写CUDA内核，增加了开发成本。因此，在模型设计阶段就应尽量避免使用高度自定义的算子，优先选择TensorRT原生支持的算子集合。

4.2 精度校准：FP16与INT8的权衡与实施

精度校准是TensorRT引擎构建中最能体现“工程经验”的环节。FP16（半精度）和INT8（8位整数）分别代表了不同的性能-精度权衡。FP16将权重和激活值从32位浮点缩减为16位，显存占用减半，计算吞吐翻倍，但动态范围缩小，可能导致梯度下溢或精度损失。INT8则更进一步，将数据量化为8位整数，需要额外的校准过程来确定量化参数（scale和zero-point），以最小化量化误差。

FP16的实施相对简单：在构建引擎时设置BuilderFlag.FP16即可。但需要注意，并非所有层都适合FP16——某些对精度敏感的层（如BatchNorm的缩放操作）可能需要保持FP32。TensorRT允许通过layer.precision属性进行逐层控制。INT8则复杂得多，其核心是校准（Calibration）：需要准备一个校准数据集（通常500-1000张代表性图片），TensorRT会运行模型并统计每层激活值的分布，然后选择最优的量化阈值。

INT8校准的完整流程包括：准备校准数据集（应覆盖真实推理时的输入分布，避免单一场景偏差）→ 创建Int8Calibrator子类并实现get_batch_size、get_batch、read_calibration_cache和write_calibration_cache方法 → 构建引擎时设置INT8标志并传入校准器。校准算法通常选择EntropyCalibrator2（熵校准），它在精度和性能间取得了较好的平衡。

一个典型的INT8校准器实现片段：

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

class EntropyCalibrator(trt.IInt8EntropyCalibrator2):
    def __init__(self, data_loader, cache_file):
        super().__init__()
        self.data_loader = data_loader
        self.cache_file = cache_file
        self.batch_size = data_loader.batch_size
        self.current_index = 0
        # 分配GPU内存
        self.d_input = cuda.mem_alloc(self.batch_size * 3 * 224 * 224 * 4)
    
    def get_batch_size(self):
        return self.batch_size
    
    def get_batch(self, names):
        if self.current_index >= len(self.data_loader):
            return None
        batch = self.data_loader[self.current_index]
        self.current_index += 1
        cuda.memcpy_htod(self.d_input, batch.astype(np.float32))
        return [int(self.d_input)]
    
    def read_calibration_cache(self):
        if os.path.exists(self.cache_file):
            return open(self.cache_file, "rb").read()
        return None
    
    def write_calibration_cache(self, cache):
        with open(self.cache_file, "wb") as f:
            f.write(cache)

在实战中，我们为一个语义分割模型实施INT8时，发现mIoU从0.72降至0.68，损失超过预期。分析后发现，模型的输出层对数值精度极其敏感。解决方案是使用trt.layers.Layer.PRECISION将输出层强制设为FP16，最终mIoU恢复至0.71，同时推理速度提升了2.3倍。这说明了逐层精度控制的重要性。

4.3 动态形状配置：灵活性与显存效率的平衡

动态形状（Dynamic Shape）是TensorRT支持输入尺寸在运行时变化的关键机制，也是多模型共享GPU资源时最核心的配置项。其原理是：在构建引擎时，为每个动态维度指定一个范围（最小值、最优值、最大值），TensorRT会基于这些范围进行内核自动调优（kernel auto-tuning）和显存池预分配。运行时，实际输入形状必须落在该范围内，且引擎会为每个不同的形状选择最优的内核实现。

配置动态形状的核心是OptimizationProfile。每个profile包含三组维度值：min（最小形状）、opt（最优形状，通常为最常见形状）、max（最大形状）。TensorRT会以opt为基准进行性能优化，因此opt应设置为推理时最频繁出现的形状。一个引擎可以包含多个profile，运行时通过context.set_optimization_profile_async切换。

构建动态形状引擎的代码流程如下：

builder = trt.Builder(logger)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, logger)
parser.parse_from_file("model.onnx")

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

profile = builder.create_optimization_profile()
# 假设输入张量名为"input"，维度为(batch, 3, H, W)
profile.set_shape("input", min=(1, 3, 512, 512), opt=(4, 3, 768, 768), max=(8, 3, 1024, 1024))
config.add_optimization_profile(profile)

engine = builder.build_serialized_network(network, config)

在运行时，需要根据实际输入形状调用context.set_input_shape，并确保输入缓冲区的尺寸足够容纳最大形状。一个常见的错误是：在创建CUDA内存时只分配了当前输入形状所需的空间，当输入形状变化时导致越界。正确的做法是始终按max形状分配显存。

在多模型共享GPU的场景下，动态形状配置需要与调度策略协同。例如，当多个模型交替执行时，每个模型的引擎应尽量使用较小的max值，以减少显存预分配，为其他模型留出空间。同时，可以考虑使用context.set_timing_cache来复用构建时的调优结果，加速后续引擎的构建。

4.4 实战案例：多模型场景下的引擎构建与调优

我们以一个实际的视频分析平台为例，该平台需要同时运行目标检测（YOLOv5s）、人脸识别（FaceNet）和姿态估计（HRNet）三个模型，共享一块NVIDIA A10 GPU（24GB显存）。初始方案是为每个模型构建固定形状的FP16引擎，输入分辨率分别为640×640、160×160和256×192。上线后发现两个问题：一是GPU显存峰值达到20GB，接近上限；二是当视频流分辨率波动时（如从1080p切换到720p），固定形状引擎无法适配，需要重新构建。

针对显存问题，我们分析了三个模型的显存占用分布：YOLOv5s约4GB，FaceNet约1.5GB，HRNet约3GB，加上TensorRT上下文和CUDA上下文，总占用约10GB。但固定形状引擎在构建时会为最大形状预分配工作空间，导致实际占用远高于模型本身。解决方案是改用动态形状，并为每个模型设置合理的min/opt/max：YOLOv5s的输入范围设为(1,3,320,320)到(1,3,1280,1280)，opt为(1,3,640,640)；FaceNet和HRNet类似。同时，将工作空间限制从默认的1GB降至512MB，通过config.set_memory_pool_limit控制。

针对分辨率波动问题，我们为每个模型配置了两个profile：一个针对1080p输入，一个针对720p输入。运行时，调度器根据当前视频流的分辨率选择对应的profile。这一改动使显存峰值降至14GB，为其他任务留出了充足空间。

复盘这一案例，我们总结了三条经验：第一，动态形状的opt值应基于真实推理频率统计，而非主观猜测；第二，显存优化不能只看模型权重大小，必须考虑TensorRT的工作空间和上下文开销；第三，多profile切换有微小的性能开销（约0.1ms），在延迟敏感场景下应尽量减少切换频率。