分割階段¶

此階段為選用，且由使用者啟用。它指示編譯器將節點分離為應在 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 ()

此分割是有效的，但分割是次優的。由於我們在圖的線性遍歷中在 Torch 和 TensorRT 目標之間交替，因此這些算術運算和 `aten::lgamma` 運算都被分割成自己的區段。

%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` 標記用於防止合併為條件節點和迴圈創建的區段，這些區段在圖縫合中作為特殊情況處理，不應與相同類型的相鄰區段合併。