分区阶段¶

此阶段是可选的，由用户启用。它指示编译器将节点分为应在 PyTorch 中运行的节点和应在 TensorRT 中运行的节点。分离的标准包括：缺少转换器，运算符被用户明确设置为在 PyTorch 中运行，或者节点具有一个标志，该标志告诉分区通过模块回退通道在 PyTorch 中运行。

从宏观层面看，Torch-TensorRT 分区阶段执行以下操作：

分割。按顺序遍历运算符集，并验证每个运算符是否有转换器。然后，粗略地将图分成 Torch-TensorRT 可以支持的部分和 Torch-TensorRT 不能支持的部分。
依赖分析。对于每个要编译的运算符，都有一个“完整依赖图”，这意味着每个输入都可以追溯到张量或张量列表形式的输入。在分割之后遍历所有段，然后进行依赖分析，以确保 TensorRT 段只有张量/张量列表输入和输出。
形状分析。对于每个段，从用户提供的输入形状开始，计算输入和输出形状。形状可以通过使用 JIT 运行图来计算。
转换。每个 TensorRT 段都将转换为 TensorRT 引擎。此部分在 compiler.cpp 中完成，但它仍然是分区过程中的一个阶段。
拼接。将所有 TensorRT 引擎与 PyTorch 节点一起拼接。

以下是这些文件中每个功能的简要说明：

PartitonInfo.h/.cpp¶

core/partitioning/PartitionInfo.h

用于分区的自动回退 API。

SegmentedBlock.h/.cpp¶

core/partitioning/SegmentedBlock.h

分割后用于维护每个段信息的主要数据结构。

shape_analysis.h/.cpp¶

core/partitioning/shape_analysis.h

通过在 JIT 中运行每个段来获取其形状的代码实现。

partitioning.h/.cpp¶

core/partitioning/partitioning.h

分区阶段的 API 和主要代码实现。

自动回退¶

要启用自动回退功能，您可以在 Python 中设置以下属性：

import torch
import torch_tensorrt as torchtrt

...
model = MyModel()
ts_model = torch.jit.script(model)
trt_model = torchtrt.ts.compile(model, **{
  ...
  "min_block_size" : 3,
  "torch_executed_ops": ["aten::add"],
  "torch_executed_modules": [],
})

enabled：默认情况下，自动回退将关闭。通过将其设置为 True 来启用。
min_block_size：必须满足被转换为 TensorRT 的连续操作的最小数量。例如，如果设置为 3，则必须有 3 个连续的受支持操作，然后此段将被转换。
forced_fallback_ops：一个字符串列表，它们将是用户明确希望在 PyTorch 节点中执行的操作的名称。

#include "torch/script.h"
#include "torch_tensorrt/torch_tensorrt.h"

...
auto in = torch::randn({1, 3, 224, 224}, {torch::kCUDA});

auto mod = torch::jit::load("trt_ts_module.ts");
auto input_sizes =  std::vector<torchtrt::InputRange>{{in.sizes()}};
torchtrt::ts::CompileSpec cfg(input_sizes);
cfg.min_block_size = 2;
cfg.torch_executed_ops.push_back("aten::relu");
auto trt_mod = torchtrt::ts::compile(mod, cfg);
auto out = trt_mod.forward({in});

依赖感知分区¶

在分割过程中，Torch-TensorRT 使用输入 TorchScript 节点的依赖图来减少创建的段数。请看 tests/core/partitioning/test_segmentation.cpp 中 Partitioning.SegmentModelWithDependencyAwareness 的示例。

graph(%x : Tensor, %y : Tensor):
    %3 : int = prim::Constant[value=0]()
    %20 : int = prim::Constant[value=1]()
    %add : Tensor = aten::add(%x, %y, %20)
    %x_lgamma : Tensor = aten::lgamma(%x)
    %mul : Tensor = aten::mul(%x, %y)
    %y_lgamma : Tensor = aten::lgamma(%y)
    %div : Tensor = aten::div(%x, %y)
    %div_lgamma : Tensor = aten::lgamma(%div)
    %27 : Tensor[] = prim::ListConstruct(%x_lgamma, %y_lgamma, %div_lgamma, %add, %mul)
    %12 : Tensor = aten::cat(%27, %3)
    return (%12)

在这个图中，aten::lgamma 不受转换支持，必须在 Torch 回退段中进行分区。如果 Torch-TensorRT 使用贪婪分割策略，即按顺序遍历输入图中的节点，并将具有相同目标（TensorRT 或 Torch）的操作收集到一个段中，直到遇到具有不同目标的操作，则结果分区将包含 7 个段，其中许多段只有一个操作。

Segment Block @0:
    Target: TensorRT

    Graph: graph(%x : Tensor,
        %y : Tensor):
    %3 : int = prim::Constant[value=1]()
    %0 : Tensor = aten::add(%x, %y, %3)
    return ()

Segment Block @1:
    Target: Torch

    Graph: graph(%x : Tensor):
    %0 : Tensor = aten::lgamma(%x)
    return ()

Segment Block @2:
    Target: TensorRT

    Graph: graph(%x : Tensor,
        %y : Tensor):
    %0 : Tensor = aten::mul(%x, %y)
    return ()

Segment Block @3:
    Target: Torch

    Graph: graph(%y : Tensor):
    %0 : Tensor = aten::lgamma(%y)
    return ()

Segment Block @4:
    Target: TensorRT

    Graph: graph(%x : Tensor,
        %y : Tensor):
    %0 : Tensor = aten::div(%x, %y)
    return ()

Segment Block @5:
    Target: Torch

    Graph: graph(%1 : Tensor):
    %0 : Tensor = aten::lgamma(%1)
    return ()

Segment Block @6:
    Target: TensorRT

    Graph: graph(%1 : Tensor,
        %2 : Tensor,
        %3 : Tensor,
        %4 : Tensor,
        %5 : Tensor):
    %7 : int = prim::Constant[value=0]()
    %0 : Tensor[] = prim::ListConstruct(%1, %2, %3, %4, %5)
    %6 : Tensor = aten::cat(%0, %7)
    return ()

这个分区是有效的，但分割不是最优的。这些算术运算和 aten::lgamma 运算在图的线性遍历中交替使用 Torch 和 TensorRT 目标时，每个都被分割成自己的段。

%add : Tensor = aten::add(%x, %y, %20)
%x_lgamma : Tensor = aten::lgamma(%x)
%mul : Tensor = aten::mul(%x, %y)
%y_lgamma : Tensor = aten::lgamma(%y)
%div : Tensor = aten::div(%x, %y)
%div_lgamma : Tensor = aten::lgamma(%div)

此段中的每个算术运算都只依赖于常量和输入 %x 和 %y。aten::lgamma 运算依赖于输入 %x、%y 以及 aten::div 的输出。这意味着我们可以在不改变图行为的情况下将输入图的这部分重写如下。这个重新排序的操作系列可以使用上述贪婪分割方法干净地分成仅 2 个段。

%add : Tensor = aten::add(%x, %y, %20)
%mul : Tensor = aten::mul(%x, %y)
%div : Tensor = aten::div(%x, %y)
%x_lgamma : Tensor = aten::lgamma(%x)
%y_lgamma : Tensor = aten::lgamma(%y)
%div_lgamma : Tensor = aten::lgamma(%div)

通过在基本的贪婪分割方法中增加对操作之间依赖关系的感知，我们可以实现相同的分区而无需重写图。现在，我们将在遍历图时同时维护 Torch 和 TensorRT 目标段。我们只会在遇到一个既依赖于段中某个操作又具有不同目标的操作时才完成一个段。这将允许分区通过跨段边界重新排序节点来创建更大的段，同时保证我们不会通过相对于其依赖关系重新排序节点来修改图的行为。在此示例中，我们将算术操作收集到 TensorRT 段中，并将 aten::lgamma 操作收集到 Torch 段中。当我们遇到 %div_lgamma : Tensor = aten::lgamma(%div) 操作时，我们可以看到它依赖于当前 TensorRT 段中的 %div : Tensor = aten::div(%x, %y)。这将触发包含 aten::div 操作的 TensorRT 段的最终确定，以确保它在最终分区中出现在其依赖项之前。包含 aten::lgamma 操作的 Torch 段在我们遇到针对 TensorRT 且依赖于 aten::lgamma 操作结果的 prim::ListConstruct 操作时最终确定。

Segment Block @0:
    Target: TensorRT

    Graph: graph(%x : Tensor,
        %y : Tensor):
    %3 : int = prim::Constant[value=1]()
    %0 : Tensor = aten::add(%x, %y, %3)
    %4 : Tensor = aten::mul(%x, %y)
    %5 : Tensor = aten::div(%x, %y)
    return ()

Segment Block @1:
    Target: Torch

    Graph: graph(%x : Tensor,
        %y : Tensor,
        %5 : Tensor):
    %0 : Tensor = aten::lgamma(%x)
    %2 : Tensor = aten::lgamma(%y)
    %4 : Tensor = aten::lgamma(%5)
    return ()

Segment Block @2:
    Target: TensorRT

    Graph: graph(%1 : Tensor,
        %2 : Tensor,
        %3 : Tensor,
        %4 : Tensor,
        %5 : Tensor):
    %7 : int = prim::Constant[value=0]()
    %0 : Tensor[] = prim::ListConstruct(%1, %2, %3, %4, %5)
    %6 : Tensor = aten::cat(%0, %7)
    return ()

在某些情况下，这种方法可能会在分区中创建具有相同目标的相邻段。作为清理步骤，我们可以合并这些相邻段，以进一步减少最终分区中的段数。合并段步骤会识别图中相邻、具有相同目标且未标记为 do_not_merge 的段列表。这些段中的节点将被组合成一个单一的新段，该新段将替换分区中合并的段。do_not_merge 标记用于防止合并为条件节点和循环创建的段，这些段在图拼接中作为特殊情况处理，不应与相同类型的相邻段合并。

Dynamo 的分层分区器¶

分层分区器是标准 TensorRT 分区器的扩展，通过考虑后端优先级和运算符支持，允许更复杂的分区策略。当您希望根据模型的能力和优先级将模型的不同部分分布到多个后端时，这尤其有用。

我们目前支持分层邻接分区器，它通过以下功能扩展了标准邻接分区器：

后端优先级排序：根据优先级顺序将运算符分配给后端，确保运算符分配给支持它们的最高优先级后端。
多后端支持：根据运算符支持在多个后端之间分配模型执行。

请参阅 hierarchical_partitioner_example 了解更多详情。