分区阶段¶

该阶段是可选的，由用户启用。它指示编译器将节点分为应在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段中的算术操作和Torch段中的aten::lgamma操作。当我们遇到%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标记用于防止合并为条件节点和循环创建的段，这些段在图缝合中作为特殊情况处理，不应与相同类型的相邻段合并。