分割槽階段¶

此階段是可選的，由使用者啟用。它指示編譯器將節點分為應在 PyTorch 中執行的節點和應在 TensorRT 中執行的節點。分離的標準包括：缺少轉換器、使用者顯式將操作設定為在 PyTorch 中執行，或者節點具有由模組回退 Pass 通知分割槽應在 PyTorch 中執行的標誌。

總的來說，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 標記用於防止合併為條件節點和迴圈建立的段，這些段在圖拼接中作為特殊情況處理，不應與相同型別的相鄰段合併。