3. TensorRT部署流程概览

在深入语义分割模型的TensorRT部署之前，我们首先需要建立一张完整的“作战地图”。TensorRT的部署并非一个孤立的步骤，而是一条从模型训练完成到生产环境高效推理的流水线。理解这条流水线的全貌，尤其是输出解析与显存管理这两个环节所处的位置，是避免“模型能跑但性能差”或“程序莫名崩溃”等问题的关键。本章将带你走完从ONNX导出到推理执行的完整旅程，并重点剖析这两个核心模块在流程中的角色与交互关系。

上图清晰地展示了TensorRT部署的三大阶段。可以看到，显存管理贯穿了引擎构建后的整个生命周期，而输出解析则是推理执行阶段紧接在GPU计算完成后的关键一步。下面我们将逐一拆解每个阶段的具体任务，并深入探讨这两个核心模块的工程实现细节。

3.1 阶段一：模型转换与引擎构建——奠定基础

3.1.1 核心原理：从PyTorch/TF到TensorRT的“翻译”

TensorRT本质上是一个深度学习推理优化器和运行时引擎。它无法直接读取PyTorch的.pt文件或TensorFlow的.pb文件。因此，第一步需要将训练好的模型转换为一种中间表示——ONNX（Open Neural Network Exchange）。ONNX充当了“通用语言”的角色，将不同框架的模型统一描述为计算图。

导出ONNX时，需要特别注意动态轴的设置。对于语义分割模型，输入尺寸通常是固定的（例如512x512），但为了灵活性，我们通常将批量大小（batch size）设为动态轴。在导出命令中，这通过dynamic_axes参数实现。

import torch
import torchvision

model = torchvision.models.segmentation.fcn_resnet50(pretrained=True)
model.eval()

# 构造一个示例输入，用于指定输入输出的形状
dummy_input = torch.randn(1, 3, 512, 512)

# 导出ONNX，设置动态批量大小
torch.onnx.export(
    model,
    dummy_input,
    "fcn_resnet50.onnx",
    opset_version=17,
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={
        "input": {0: "batch_size"},
        "output": {0: "batch_size"}
    }
)
print("ONNX模型导出成功！")

3.1.2 操作流程：使用trtexec构建引擎

获得ONNX模型后，我们使用TensorRT自带的trtexec工具或编写Python/C++代码来构建优化后的引擎。构建过程涉及层融合、精度校准（如INT8量化）、内核自动调优等复杂操作。构建出的引擎文件（.engine或.plan）是序列化的，它包含了针对特定GPU架构和TensorRT版本的优化指令。

使用命令行构建引擎的命令示例如下：

trtexec --onnx=fcn_resnet50.onnx \
        --saveEngine=fcn_resnet50.engine \
        --fp16 \
        --minShapes=input:1x3x512x512 \
        --optShapes=input:4x3x512x512 \
        --maxShapes=input:8x3x512x512

3.2 阶段二：运行时初始化——加载与准备

3.2.1 核心原理：反序列化与上下文创建

在推理程序启动时，我们需要将之前构建好的.engine文件加载到内存中。这个过程称为反序列化。TensorRT运行时（IRuntime）负责解析引擎文件，并创建一个可执行的ICudaEngine对象。随后，我们需要基于该引擎创建一个或多个IExecutionContext。上下文（Context）是实际执行推理的“工作区”，它持有中间激活值的内存、内核执行流等状态。

这里有一个关键点：一个引擎可以创建多个上下文，用于多流并发推理。但每个上下文都会占用额外的显存，因此需要权衡并发度与显存容量。

3.2.2 操作流程：显存管理的起点

在创建上下文后，我们需要为模型的输入和输出张量分配GPU显存。这是显存管理的第一个显式操作点。TensorRT提供了灵活的显存分配方式，推荐使用IExecutionContext.setTensorAddress方法直接绑定预先分配的显存地址。

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

# 加载引擎
runtime = trt.Runtime(trt.Logger(trt.Logger.WARNING))
with open("fcn_resnet50.engine", "rb") as f:
    engine = runtime.deserialize_cuda_engine(f.read())

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

# 获取输入输出张量名称
input_name = "input"
output_name = "output"

# 设置动态形状的实际输入尺寸
context.set_input_shape(input_name, (1, 3, 512, 512))

# 获取绑定后的显存大小
input_size = trt.volume(context.get_tensor_shape(input_name)) * trt.float32.itemsize
output_size = trt.volume(context.get_tensor_shape(output_name)) * trt.float32.itemsize

# 分配GPU显存
d_input = cuda.mem_alloc(input_size)
d_output = cuda.mem_alloc(output_size)

# 绑定显存地址
context.set_tensor_address(input_name, int(d_input))
context.set_tensor_address(output_name, int(d_output))

print(f"输入张量大小: {input_size} bytes, 输出张量大小: {output_size} bytes")

3.3 阶段三：推理执行与输出解析——核心战场

3.3.1 核心原理：异步执行与数据搬运

推理执行的核心是调用context.execute_async_v3方法。该方法会启动GPU内核，并立即返回，不会阻塞CPU。为了等待GPU计算完成，我们需要使用CUDA事件（Event）或流（Stream）进行同步。典型的推理循环包括：

1. 预处理：将输入图像从CPU拷贝到GPU（cuda.memcpy_htod_async）。
2. 执行推理：调用execute_async_v3。
3. 同步等待：通过cuda.Event或stream.synchronize()等待GPU完成。
4. 输出解析：将输出张量从GPU拷贝回CPU（cuda.memcpy_dtoh_async），并进行后处理。

3.3.2 输出解析：从原始数据到可视化结果

对于语义分割模型，TensorRT的输出通常是一个形状为[batch, num_classes, H, W]的浮点张量。这个张量代表每个像素属于各个类别的置信度分数（Logits）。要得到最终的分割掩码，我们需要执行以下步骤：

• Softmax归一化（可选）：将分数转换为概率分布。
• Argmax操作：在类别维度上取最大值索引，得到每个像素的类别ID。
• 调色板映射：将类别ID映射到RGB颜色，生成可视化图像。

以下是一个典型的输出解析代码片段：

import numpy as np
import cv2

# 假设 output_data 是从GPU拷贝回来的numpy数组，形状为 (1, 21, 512, 512)
output_data = np.zeros((1, 21, 512, 512), dtype=np.float32)
cuda.memcpy_dtoh(output_data, d_output)

# 移除batch维度，得到 (21, 512, 512)
scores = output_data[0]

# 执行argmax，得到 (512, 512) 的类别索引
class_ids = np.argmax(scores, axis=0).astype(np.uint8)

# 定义调色板（PASCAL VOC数据集）
palette = np.array([
    (0, 0, 0), (128, 0, 0), (0, 128, 0), (128, 128, 0),
    (0, 0, 128), (128, 0, 128), (0, 128, 128), (128, 128, 128),
    (64, 0, 0), (192, 0, 0), (64, 128, 0), (192, 128, 0),
    (64, 0, 128), (192, 0, 128), (64, 128, 128), (192, 128, 128),
    (0, 64, 0), (128, 64, 0), (0, 192, 0), (128, 192, 0),
    (0, 64, 128)
], dtype=np.uint8)

# 映射到彩色图像
segmentation_map = palette[class_ids.reshape(-1)].reshape(512, 512, 3)
segmentation_map = cv2.cvtColor(segmentation_map, cv2.COLOR_RGB2BGR)

# 保存结果
cv2.imwrite("segmentation_result.png", segmentation_map)

3.3.3 案例拆解：动态形状下的显存与输出管理

假设我们部署的语义分割服务需要处理不同分辨率的输入图像（如来自不同摄像头的视频流）。我们构建引擎时设置了动态形状，允许输入在[1x3x256x256]到[8x3x1024x1024]之间变化。

初始现象与目标：服务启动后，发现当输入分辨率从512x512切换到1024x1024时，程序抛出std::bad_alloc异常，或者推理结果出现乱码。

分析判断过程：

• 首先，检查显存分配逻辑。发现代码中在初始化时只分配了基于512x512尺寸的缓冲区。
• 当输入尺寸增大时，set_input_shape被调用，但输出缓冲区的大小并未更新，导致GPU写入越界。
• 其次，输出解析代码中，output_data数组的大小是硬编码的(1, 21, 512, 512)，导致从GPU拷贝的数据量不匹配。

实际操作步骤：

1. 动态查询形状：在每次推理前，调用context.get_tensor_shape(output_name)获取当前实际的输出维度。
2. 条件性重新分配：如果当前输出尺寸所需的显存大于已分配的缓冲区，则释放旧缓冲区并重新分配。
3. 同步更新CPU端数组：根据查询到的形状，重新创建output_data numpy数组。

结果对比与复盘：修复后，服务可以稳定处理动态分辨率输入。此案例揭示了显存管理与输出解析之间的强耦合关系：输出解析的代码逻辑必须基于运行时实际的张量形状，而显存分配必须为最坏情况预留空间或支持动态扩容。

3.4 常见误区与总结复盘

总结来看，TensorRT部署流程是一个环环相扣的工程链路。从ONNX导出时的算子兼容性检查，到引擎构建时的精度与速度权衡，再到运行时显存的高效复用和输出数据的正确解析，每一步都值得仔细推敲。对于语义分割任务，尤其要关注输出张量的维度变换和数据类型转换，确保从GPU到CPU的数据流是准确且高效的。掌握本章的流程概览，你就为后续深入优化打下了坚实的基础。