分区阶段¶

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

在较高层次上，Torch-TensorRT 分区阶段执行以下操作

分段。按顺序遍历运算符集合，并验证每个运算符是否存在转换器。然后，大致将图分成 Torch-TensorRT 可以支持的部分和 Torch-TensorRT 无法支持的部分。
依赖性分析。对于每个要编译的运算符，都存在“完整的依赖关系图”，这意味着每个输入都可以追溯到作为 Tensor 或 TensorList 的输入。遍历分段后的所有段，然后进行依赖性分析，以确保 TensorRT 段仅具有 Tensor/TensorList 输入和输出。
形状分析。对于每个段，从用户提供的输入形状开始，计算输入和输出形状。可以通过使用 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 段的最终确定，以保证它将出现在其依赖项之前在最终分区中。当我们遇到以 TensorRT 为目标并且依赖于 aten::lgamma 操作结果的 prim::ListConstruct 操作时，包含 aten::lgamma 操作的 Torch 段将被最终确定。

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 标记用于防止合并为条件节点和循环创建的段，这些段在图拼接中作为特殊情况处理，不应与相同类型的相邻段合并。