分块阶段¶

此阶段是可选的，由用户启用。它指示编译器将节点分离为应在 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 段的完成，以保证它在最终块中出现在其依赖项之前。包含 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 标记用于防止合并为条件节点和循环创建的段，这些段在图拼接中作为特殊情况处理，不应与相同类型的相邻段合并。