捷徑

降低階段

降低階段由多個 pass 組成,這些 pass 是將圖形從高階表示法映射到低階表示法的運算。每個 pass 都執行特定的操作,例如內聯方法呼叫。其想法是顯著減少轉換階段在實際映射到 TensorRT 時需要處理的內容。我們的目標是更接近 1->1 運算子轉換,而不是尋找適用的子圖,限制轉換器的數量並縮小每個轉換器的範圍。

您可以透過將日誌等級設定為 Level::kGraph 來查看每個 pass 的效果

使用的 Pass

EliminateCommonSubexpression

移除圖形中的常見子運算式

消除無效程式碼

無效程式碼消除將檢查節點是否具有副作用,如果具有副作用則不會刪除它。

消除例外或 Pass 模式

腳本化模組中的常見模式是維度保護,如果輸入維度不是預期的維度,則會拋出例外。

%1013 : bool = aten::ne(%1012, %24) # ~/.local/lib/python3.6/site-packages/torch/nn/modules/batchnorm.py:248:11
    = prim::If(%1013) # ~/.local/lib/python3.6/site-packages/torch/nn/modules/batchnorm.py:248:8
    block0():
        = prim::RaiseException(%23) # ~/.local/lib/python3.6/site-packages/torch/nn/modules/batchnorm.py:249:12
    -> ()
    block1():
    -> ()

由於我們在編譯時解析所有這些,並且 TensorRT 圖形中沒有例外,因此我們只需將其移除。

消除冗餘保護

消除運算子的冗餘保護,這些運算子的輸出完全由其輸入決定,即如果此類運算子的輸入受到保護,我們允許移除運算子輸出的保護

凍結模組

凍結屬性並內聯常數和模組。傳播圖形中的常數。

融合 AddMM 分支

腳本化模組中的常見模式是不同維度的張量使用不同的建構來實作線性層。我們將這些不同的變體融合為單一變體,該變體將被 Unpack AddMM pass 捕獲。

%ret : Tensor = prim::If(%622)
block0():
  %ret.1 : Tensor = aten::addmm(%self.fc.bias, %x9.1, %3677, %3, %3)
  -> (%ret.1)
block1():
  %output.1 : Tensor = aten::matmul(%x9.1, %3677)
  %output0.1 : Tensor = aten::add_(%output.1, %self.fc.bias, %3)
  -> (%output0.1)

我們將這組區塊融合到如下圖形中

%ret : Tensor = aten::addmm(%self.fc.bias, %x9.1, %3677, %3, %3)

融合線性

比對 aten::linear 模式並將其融合為單一 aten::linear。此 pass 將 JIT 產生的 addmm 或 matmul + add 融合回線性

融合 Flatten Linear

TensorRT 隱式地將輸入層展平為高於 1D 的全連接層。因此,當存在 aten::flatten -> aten::linear 模式時,我們會移除 aten::flatten

降低圖形

給定一個方法圖形,其第一個引數為 %self,將其降低為一個圖形,其中所有屬性存取都替換為圖形的明確輸入(而不是在 %self 上執行的 prim::GetAttr 的結果)。傳回一個元組 (圖形, 參數),其中圖形的最後 module.parameters.size() 個輸入是此方法中使用的可訓練參數。剩餘的輸入是函數的真實輸入。

降低元組

  • LowerSimpleTuples:

移除 TupleConstruct 和 TupleUnpack 比對的元組,但在 if 語句、迴圈以及作為輸入/輸出時保留元組

  • LowerAllTuples:

移除_所有_元組,如果某些元組無法移除,則會引發錯誤,ONNX 使用此功能來確保轉換前沒有元組,但對於輸入包含元組的圖形將不起作用。

模組回退

模組回退由兩個降低 pass 組成,必須成對執行。第一個 pass 在凍結之前執行,以在圖形中應在 PyTorch 中執行的模組周圍放置分隔符。第二個 pass 在凍結後標記這些分隔符之間的節點,以表示它們應在 PyTorch 中執行。

  • NotateModuleForFallback

在凍結之前在模組呼叫周圍放置分隔節點,以表示圖形中節點應在 PyTorch 中執行的位置

  • MarkNodesForFallback

尋找分隔符,然後標記分隔符之間的所有節點,以告知分割在 PyTorch 中執行它們

Peephole Optimize

此最佳化 pass 的目的是捕獲您可能感興趣執行的所有小的、易於捕獲的 Peephole 最佳化。

目前,它執行以下操作
  • 消除無運算的 'expand' 節點

  • 將 x.t().t() 簡化為 x

移除 Contiguous

移除 contiguous 運算子,因為我們正在執行 TensorRT,記憶體已經是連續的。

移除 Dropout

移除 dropout 運算子,因為我們正在執行推論。

移除 To

移除執行型別轉換的 aten::to 運算子,因為 TensorRT 會自行管理。重要的是,這是最後執行的 pass 之一,以便其他 pass 有機會將所需的型別轉換運算子移出主要命名空間。

Unpack AddMM

aten::addmm 解壓縮為 aten::matmulaten::add_(以及額外的 trt::const 運算子,以凍結 TensorRT 圖形中的偏差)。這讓我們可以重複使用 aten::matmulaten::add_ 轉換器,而無需專用的轉換器。

Unpack LogSoftmax

aten::logsoftmax 解壓縮為 aten::softmaxaten::log。這讓我們可以重複使用 aten::softmaxaten::log 轉換器,而無需專用的轉換器。

Unroll Loops

展開相容迴圈(例如,足夠短)的運算,以便您只需執行一次迴圈。

將 Tile 替換為 Repeat

移除 dropout 運算子,因為我們正在執行推論。

文件

存取 PyTorch 的完整開發者文件

檢視文件

教學

取得針對初學者和進階開發者的深入教學

檢視教學

資源

尋找開發資源並獲得問題解答

檢視資源