📚 课程目录第2章 TensorRT构建动态Shape引擎的核心流程

2. TensorRT构建动态Shape引擎的核心流程

动态Shape引擎的构建并不是在TensorRT里“打开一个开关”，而是要把动态信息从ONNX模型一路传递到TensorRT的运行时。整条链路可以概括为：ONNX导出时用dynamic_axes声明哪些轴可变，TensorRT解析ONNX时把这些轴识别为-1，随后通过OptimizationProfile为每个可变轴给出最小、最优、最大三组形状，最后在构建配置中绑定Profile并设置工作空间。任何一环缺失，都会导致引擎退化为静态Shape，或在运行时因形状越界而报错。

一、ONNX侧：用dynamic_axes声明可变轴

动态Shape的源头在ONNX。以PyTorch导出为例，假设模型输入为input，形状是[batch, 3, height, width]，输出为output，形状是[batch, num_classes]。如果希望batch、height、width都可变，导出时应显式声明：

torch.onnx.export(
    model, dummy_input, "model_dynamic.onnx",
    input_names=["input"], output_names=["output"],
    dynamic_axes={
        "input":  {0: "batch", 2: "height", 3: "width"},
        "output": {0: "batch"}
    },
    opset_version=13
)

这里dynamic_axes的键是张量名，值是“轴索引→轴名称”的映射。轴名称只是可读标签，真正起作用的是轴索引。导出后可以用Netron查看，动态轴会显示为batch、height、width，而不是固定数字。需要注意：如果只声明了输入动态轴，而输出形状依赖这些轴，输出轴也应一并声明，否则TensorRT可能推断出错误的输出维度。此外，某些算子对动态轴支持有限，例如带固定kernel_size的池化在H/W动态时通常没问题，但涉及reshape、transpose时若硬编码了维度，就会在动态场景下失败。

二、Builder与Network：解析模型并识别动态输入

进入TensorRT后，第一步是创建IBuilder和INetworkDefinition，再用ONNX解析器把模型读入网络：

IBuilder* builder = createInferBuilder(gLogger);
INetworkDefinition* network = builder->createNetworkV2(
    1U << static_cast<int>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH));
auto parser = nvonnxparser::createParser(*network, gLogger);
parser->parseFromFile("model_dynamic.onnx",
    static_cast<int>(ILogger::Severity::kWARNING));

解析完成后，应遍历network->getInput(i)和getOutput(i)，检查每个张量的getDimensions()。动态轴在TensorRT中表现为-1。例如输入可能显示为[-1, 3, -1, -1]，输出为[-1, num_classes]。如果这里看到的是固定数字，说明ONNX导出时没有正确声明动态轴，需要回到上一步修正。这一步是判断“动态信息是否成功传递”的关键检查点。

三、OptimizationProfile：为每个动态轴给出形状范围

TensorRT需要知道动态轴的实际取值范围，才能做kernel选择、显存分配和tactic调优。这由IOptimizationProfile承担。一个Profile包含每个动态输入的最小（min）、最优（opt）、最大（max）三组形状：

IOptimizationProfile* profile = builder->createOptimizationProfile();
profile->setDimensions("input", OptProfileSelector::kMIN, Dims4(1, 3, 224, 224));
profile->setDimensions("input", OptProfileSelector::kOPT, Dims4(8, 3, 512, 512));
profile->setDimensions("input", OptProfileSelector::kMAX, Dims4(16, 3, 1024, 1024));
config->addOptimizationProfile(profile);

三组形状的含义不同：kMIN和kMAX定义了运行时允许的形状闭区间，超出即报错；kOPT是构建时优先优化的形状，TensorRT会围绕它选择性能最好的kernel。因此kOPT应设为实际业务中最常出现的形状，而不是简单取中间值。例如推理服务中batch=8、512×512最常见，就把kOPT设为该值。若kMIN与kMAX跨度过大，构建时间会显著增加，因为TensorRT要为整个区间准备tactic。

min/opt/max必须满足：对每个动态轴，min ≤ opt ≤ max。若某个轴在min和max中取值相同，则该轴实际退化为静态轴。所有动态输入都必须被Profile覆盖，否则构建时报“no profile”错误。

四、绑定Profile并构建引擎

Profile创建后要加入IBuilderConfig，再调用buildEngineWithConfig：

config->setMemoryPoolLimit(MemoryPoolType::kWORKSPACE, 1ULL << 30);
IHostMemory* serializedModel = builder->buildSerializedNetwork(*network, *config);

构建成功后，引擎中会保留Profile信息。运行时必须显式设置当前使用的Profile索引，并再次设置实际输入形状，二者必须落在该Profile的min/max范围内：

context->setOptimizationProfileAsync(0, stream);
context->setInputShape("input", Dims4(4, 3, 640, 640));
context->setTensorAddress("input", d_input);
context->setTensorAddress("output", d_output);
context->enqueueV3(stream);

如果运行时形状与Profile不匹配，setInputShape会失败或enqueueV3返回错误。多Profile场景下，可以创建多个Profile分别覆盖不同形状区间，运行时按需切换，但同一时刻只有一个Profile生效。

五、多Profile的组织方式

当业务形状分布差异很大时，单个Profile的min/max区间过宽会导致构建慢、显存高。此时可以创建多个Profile，每个覆盖一个较窄的形状区间。例如Profile 0覆盖小图（224~384），Profile 1覆盖中图（512~768），Profile 2覆盖大图（896~1024）。构建时依次addOptimizationProfile，运行时根据实际输入选择对应索引。需要注意的是，每个Profile都必须为所有动态输入设置完整的min/opt/max，不能只设置部分输入。多Profile会增加构建时间和序列化体积，但能显著改善各区间内的推理性能。

动态Shape引擎一旦序列化，Profile的min/max范围就固定了。如果后续业务出现超出范围的新形状，必须重新构建引擎，不能通过运行时修改Profile来扩展范围。

把上述流程串起来看，动态Shape构建的本质是“形状契约”的逐层传递：ONNX用dynamic_axes声明哪些轴可变，TensorRT用-1识别这些轴，Profile用min/opt/max把可变范围量化，运行时再用setInputShape把具体形状填入契约。每一层都必须一致，任何一层缺失或越界，都会在构建或推理阶段暴露为错误。掌握这条链路后，排查动态Shape问题就有了明确的定位顺序：先看ONNX维度，再看Profile覆盖，最后看运行时形状是否落在区间内。