2、TensorRT基础回顾：TensorRT推理流程、ONNX模型转换、Engine构建与序列化

在进入图像分类部署的具体实践之前，有必要先厘清TensorRT的核心工作机理。图像分类任务看似简单——输入一张图片，输出一个类别标签——但其背后涉及从训练框架到推理引擎的完整链路转换。许多初学者在部署时遇到的“结果不对”“速度没提升”“显存暴涨”等问题，根源往往不在模型本身，而是对TensorRT的推理流程、ONNX转换的细节以及Engine的生命周期管理缺乏系统性认知。本章将围绕这三条主线展开，为后续Batch策略、Layout优化和预处理一致性打下坚实的工程基础。

上图展示了TensorRT部署的完整生命周期。可以看到，整个流程被清晰地划分为三个阶段：离线阶段的模型转换、构建阶段的Engine生成、运行时阶段的推理执行。理解每个阶段的任务边界和产物形态，是避免部署混乱的第一步。下面我们逐一深入。

2.1 TensorRT推理流程：从张量到引擎的完整链路

核心原理：什么是TensorRT推理

TensorRT推理的本质，是将训练好的深度学习模型转化为针对特定GPU架构优化的计算图，并在运行时高效执行。与PyTorch或TensorFlow的即时执行模式不同，TensorRT采用“先构建、后执行”的两段式设计。构建阶段会分析网络结构，进行层融合（如将Conv+BN+ReLU合并为一个内核）、精度校准（FP16/INT8）、内核自动选择等优化；执行阶段则通过ExecutionContext提交推理任务，整个过程不再有Python解释器的开销，也不涉及动态图调度。

对于图像分类任务，推理流程可以拆解为以下五个关键步骤：

1. 输入绑定：将预处理后的图像数据（通常是NCHW或NHWC布局的浮点张量）拷贝到GPU显存中的指定地址。
2. 张量绑定：通过executeV2或enqueueV3接口，将输入输出张量的指针传递给ExecutionContext。
3. 内核执行：GPU按照构建时生成的优化计划逐层执行计算，中间结果保留在显存中。
4. 输出提取：从输出张量中读取分类得分（logits）或概率分布。
5. 后处理：执行Softmax（若需要）、TopK排序、类别索引映射等操作，得到最终分类标签。

判断方法：如何确认推理流程正确

在实际部署中，判断推理流程是否正确的黄金标准是数值一致性对比。具体做法是：取同一张测试图片，分别在原始框架（如PyTorch）和TensorRT中执行前向推理，对比输出张量的差异。对于FP32精度，两者输出的最大误差通常应小于1e-3；对于FP16，误差可能放大到1e-2量级，但仍不应影响Top-1分类结果。如果发现显著差异，应优先检查预处理是否一致（包括归一化参数、通道顺序、缩放比例），其次检查ONNX导出时是否有算子被错误映射。

案例拆解：一个典型的图像分类推理调用

假设我们有一个ResNet-50的TensorRT引擎，输入为[1, 3, 224, 224]的FP32张量。以下是一个简化的C++推理流程：

// 1. 创建运行时和反序列化引擎
IRuntime* runtime = createInferRuntime(logger);
ICudaEngine* engine = runtime->deserializeCudaEngine(engineData, engineSize);
IExecutionContext* context = engine->createExecutionContext();

// 2. 分配显存
void* inputBuffer;
void* outputBuffer;
cudaMalloc(&inputBuffer, batchSize * 3 * 224 * 224 * sizeof(float));
cudaMalloc(&outputBuffer, batchSize * 1000 * sizeof(float));

// 3. 将预处理后的图像数据拷贝到显存
cudaMemcpy(inputBuffer, hostInputData, inputSize, cudaMemcpyHostToDevice);

// 4. 绑定张量并执行推理
context->setTensorAddress("input", inputBuffer);
context->setTensorAddress("output", outputBuffer);
context->enqueueV3(stream);

// 5. 将结果拷回主机
cudaMemcpy(hostOutputData, outputBuffer, outputSize, cudaMemcpyDeviceToHost);
cudaStreamSynchronize(stream);

这段代码展示了推理的核心骨架。值得注意的是，setTensorAddress使用的是张量名称而非索引，这是TensorRT 8.5+推荐的方式，更直观且不易出错。在实际项目中，还需要处理CUDA流的同步、错误检查以及多线程环境下的ExecutionContext复用问题。

常见误区与总结

总结来说，TensorRT推理流程的核心在于“构建一次、执行多次”。理解这个理念，才能在设计系统时合理分配资源，避免在运行时阶段做昂贵的构建操作。

2.2 ONNX模型转换：从训练框架到中间表示的桥梁

核心原理：为什么需要ONNX

ONNX（Open Neural Network Exchange）是一种开放的模型表示格式，旨在打通不同深度学习框架之间的壁垒。在TensorRT部署链路中，ONNX扮演着“通用语”的角色：PyTorch模型导出为ONNX，TensorRT通过ONNX解析器将其转换为内部网络定义。选择ONNX而非直接使用PyTorch导出，是因为TensorRT对ONNX的算子覆盖更全面，且ONNX提供了清晰的静态计算图，便于优化器进行分析。

对于图像分类模型，ONNX转换的核心是算子映射。PyTorch中的nn.Conv2d、nn.BatchNorm2d、nn.ReLU等模块，在ONNX中都有对应的标准算子。但某些自定义操作（如特殊的注意力机制、动态形状处理）可能没有直接映射，需要额外处理。

操作流程：PyTorch导出ONNX的标准步骤

以下是一个标准的PyTorch到ONNX的导出流程，以ResNet-50为例：

import torch
import torchvision.models as models

# 1. 加载模型并设为评估模式
model = models.resnet50(pretrained=True)
model.eval()

# 2. 构造示例输入（与训练时的预处理一致）
dummy_input = torch.randn(1, 3, 224, 224)

# 3. 导出ONNX
torch.onnx.export(
    model,
    dummy_input,
    "resnet50.onnx",
    export_params=True,
    opset_version=17,
    do_constant_folding=True,
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}}
)

这里有几个关键参数需要特别注意：

• opset_version：ONNX算子集的版本。版本过低可能缺少某些算子，过高可能导致TensorRT解析器不支持。建议使用TensorRT文档中推荐的版本（通常为13-17）。
• dynamic_axes：声明动态维度。对于图像分类，通常只将batch维度设为动态，而保持H、W固定，以简化TensorRT的优化。
• do_constant_folding：启用常量折叠，将一些在推理时不变的运算（如BN的均值和方差）预先计算，减小模型体积。

案例拆解：ONNX转换中的典型问题

假设我们导出了一个包含nn.Upsample操作的模型（如分割模型），在TensorRT中可能遇到“Unsupported Operator”错误。此时需要检查ONNX中该算子被映射为什么形式。PyTorch的Upsample在ONNX中可能被表示为Resize算子，而TensorRT对Resize的coordinate_transformation_mode参数有特定要求。解决方案是：在导出时显式指定antialias=False，或使用torch.nn.functional.interpolate的特定模式。

另一个常见问题是动态形状导致的性能退化。如果ONNX中所有维度都是动态的，TensorRT无法进行有效的内存规划和内核选择，推理速度可能下降30%-50%。因此，对于图像分类这种输入尺寸固定的任务，建议将H、W维度固定，只保留batch维度动态。

常见误区与总结

ONNX转换是TensorRT部署中最容易出错的环节之一。掌握算子映射规则、合理设置动态维度、善用模型简化工具，可以大幅减少后续Engine构建时的排障成本。

2.3 Engine构建与序列化：从ONNX到可部署的优化产物

核心原理：Engine构建的优化机制

TensorRT Engine构建的核心目标，是针对特定GPU架构生成最优的执行计划。这个过程包括：层融合（将Conv+BN+ReLU合并为CBR结构）、精度选择（FP32/FP16/INT8）、内核自动调优（从数百个候选内核中选择最快的）、显存规划（复用中间张量的显存空间）。构建过程通常需要几秒到几分钟不等，取决于模型复杂度和GPU性能。

Engine与GPU硬件绑定，这是TensorRT的一个重要特性。在A100上构建的Engine，不能直接部署到V100上，因为内核优化和指令集不同。因此，序列化后的Engine文件必须与目标GPU型号严格匹配。

操作流程：Engine构建与序列化的标准步骤

以下是一个Python API的Engine构建流程：

import tensorrt as trt

logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
network = builder.create_network(1 << trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
parser = trt.OnnxParser(network, logger)

# 1. 解析ONNX模型
with open("resnet50.onnx", "rb") as f:
    parser.parse(f.read())

# 2. 配置构建选项
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB工作空间
config.set_flag(trt.BuilderFlag.FP16)  # 启用FP16

# 3. 构建Engine
serialized_engine = builder.build_serialized_network(network, config)

# 4. 序列化保存
with open("resnet50.engine", "wb") as f:
    f.write(serialized_engine)

这段代码展示了从ONNX到Engine的完整过程。其中EXPLICIT_BATCH标志是必须的，它允许网络定义中显式包含batch维度，这是动态batch的基础。工作空间（workspace）的大小直接影响构建时的内存使用和优化效果，但过大的工作空间可能导致显存溢出。

案例拆解：FP16与INT8的精度-性能权衡

以ResNet-50为例，在NVIDIA A10 GPU上，我们对比了不同精度配置的表现：

从表中可以看到，FP16相比FP32在延迟上几乎减半，而准确率损失可以忽略不计。INT8进一步提升了性能，但准确率下降了约0.3个百分点，且需要额外的校准数据集。对于大多数图像分类场景，FP16是性价比最高的选择；如果对延迟有极致要求且准确率容忍度较高，可以考虑INT8。

常见误区与总结

Engine构建与序列化是TensorRT部署的“最后一公里”。理解构建过程的优化机制、合理选择精度配置、遵循“离线构建、在线加载”的原则，是确保部署系统稳定高效的关键。