大綱

Torch TensorRT介紹

JIT編譯與AOT編譯方法

兩種方法的異同點

詳細要點

1. Torch TensorRT介紹

Torch TensorRT是一個優(yōu)化PyTorch模型推理性能的工具

它結(jié)合了PyTorch和NVIDIA的TensorRT

2. 兩種編譯方法

JIT編譯:靈活,支持動態(tài)圖和Python代碼

AOT編譯:支持序列化,固定輸入shape

3. 編譯流程

圖轉(zhuǎn)換、優(yōu)化、劃分、TensorRT轉(zhuǎn)換

獲得高性能優(yōu)化模型

4. JIT編譯細節(jié)

通過torch.compile調(diào)用

支持動態(tài)輸入和條件判斷

5. AOT編譯細節(jié)

通過trace和compile API實現(xiàn)

使用inputAPI支持動態(tài)shape

支持序列化模型

6. 兩種方法的異同

核心是基于同一圖優(yōu)化機制

JIT支持動態(tài),AOT支持序列化

大家好，我叫喬治。嗨，我是迪拉杰，我們都是NVIDIA的深度學(xué)習(xí)軟件工程師。今天我們在這里討論使用Torch TensorRT加速PyTorch推斷。首先，我們會給大家簡短介紹一下Torch TensorRT是什么，然后喬治將深入介紹我們優(yōu)化PyTorch模型的用戶工作流程。最后，我們將比較這兩種方法，并討論一些正在進行的未來工作。現(xiàn)在我將把話筒交給喬治。

那么，什么是Torch TensorRT呢？Torch是我們大家聚在一起的原因，它是一個端到端的機器學(xué)習(xí)框架。而TensorRT則是NVIDIA的高性能深度學(xué)習(xí)推理軟件工具包。Torch TensorRT就是這兩者的結(jié)合。我們所做的是以一種有效且易于使用的方式將這兩個框架結(jié)合起來，可以適用于各種用例和模型。Torch TensorRT是優(yōu)化PyTorch模型的一種路徑。

我將在今天的用例中分享一些方法，包括Model performance和core implementation。這意味著，如果你有一個單一的模型，并且你用我們的兩種方法之一進行優(yōu)化，你可以得到相似的性能和核心軟件實現(xiàn)。現(xiàn)在你可能會問，為什么我要選擇其中之一？因為不同的部署場景需要不同的方法。

今天我們將介紹的兩種方法是Flexible JIT方法和Serializable AOT方法。如果您的模型具有復(fù)雜邏輯，需要即時編譯，并且可能需要Python嚴格部署，那么靈活的JIT方法可能對您來說是最好的選擇。如果您需要整個圖形捕獲，需要對模型進行一些序列化，或者進行C++部署，那么AOT方法可能更適合您的用例。所以考慮到這一點，讓我們走一遍這兩個用戶流程共享的內(nèi)部路徑。

為了優(yōu)化PyTorch模型，從模型創(chuàng)建開始一直到優(yōu)化，所以我們只需要用戶提供一個預(yù)訓(xùn)練模型和一些編譯設(shè)置。這些傳遞給編譯器，即Torch TensorRT?，F(xiàn)在，在這個綠色虛線框中的所有操作都是在幕后完成的，由編譯器來處理。但是這有助于理解我們?nèi)绾潍@得性能提升。

以下是一般的內(nèi)部構(gòu)造：

我們使用ATEN trace將graph轉(zhuǎn)換過來。實際上，這意味著我們將torch操作符轉(zhuǎn)換成ATEN表示，這只是以一種稍微容易處理的方式來表示相同的操作符。之后，我們進行l(wèi)owering處理，包括常數(shù)折疊和融合以提高性能。然后我們進入劃分階段partitioning。Torch TensorRT會選擇運行哪些操作，哪些操作在Torch中運行，從而生成您在右側(cè)看到的分段圖形。最后是轉(zhuǎn)換階段，對于在右側(cè)看到的每個TensorRT圖形，他們從其ATEN操作轉(zhuǎn)換為等效的TensorRT layer，最后得到優(yōu)化后的模型。

所以在談到這個一般方法后，即針對兩種用例的一般路徑,我們將稍微深入了解JIT工作流程.,Dheeraj將討論提前工作流程.,JIT方法使您享受到了Torch.compile的好處.,其中包括復(fù)雜的Python代碼處理、自動圖形分割.,.,NVIDIA的TensorRT的優(yōu)化能力.,層張量融合、內(nèi)核自動調(diào)優(yōu)，以及您選擇層精度的能力.,

在該方法的瀏覽版中，我們能夠在不到一秒鐘的時間內(nèi)對穩(wěn)定擴散的文本到圖像進行基準測試.（4090 16fp 50epoch),那么，JIT方法背后的真正原理是什么?,嗯，從一開始就開始吧.,然后用戶只需要在該模型上調(diào)用torch.compile,并指定TensorRT作為后端.,現(xiàn)在從這里開始的一切都是在幕后進行的,.,但有助于解釋正在發(fā)生的事情.

因此，在這之后，Torch.compile將會將您的模型代碼進行拆分，然后是一個dynamo guard（什么是 dynamo guard？請參考：TorchDynamo 源碼剖析 04 - Guard, Cache, Executionhttps://zhuanlan.zhihu.com/p/630722214[1] ，可以簡單理解為dynamo后到python代碼的中間層），再然后是另一個graph。之所以會出現(xiàn)這個guard，是因為在代碼中有一個條件語句。實際上，大多數(shù)機器學(xué)習(xí)模型的跟蹤器都無法處理這樣的條件語句，因為這個條件取決于輸入的值。但是Torch Compile能夠處理它，因為這個guard在Python中運行，并且會在條件的值發(fā)生變化時觸發(fā)自動圖形重新編譯。

此外，還需注意其他因素也可能導(dǎo)致重新編譯，例如如果您傳入了不同形狀的張量或者改變了其他參數(shù)。因此，這個graph / guard / graph構(gòu)造現(xiàn)在進入了Torch Tensorrt后端，它將按照之前幻燈片中展示的相同過程進行處理。右上角的每個圖形都會被轉(zhuǎn)換為右下角您看到的Torch TensorRT優(yōu)化版本。有效地將其分成TensorRT組件和Torch組件。需要注意的關(guān)鍵是Dynamo Guard保持完好。因為他提供了圖分區(qū)的強大支持。第一級是在復(fù)雜Python代碼的Python級別上。第二級是運算符級別上分區(qū)，在TensorRT中可以進一步加速的運算符以及可以在Torch中加速的其他運算符。

總結(jié)用法，用戶只需對模型調(diào)用torch compile，指定后端tensorRT，并可選擇一些選項，然后傳遞一些輸入，它將實時進行編譯。這種編譯方法對用戶來說非?？啥ㄖ?。您可以通過精度關(guān)鍵字參數(shù)選擇層精度。您可以指定在TensorRT引擎塊中所需的最小運算符數(shù)量，等等。

就這些，接下來就交給Dheeraj討論AOT方法?，F(xiàn)在讓我們來看看Torch TensorRT的AOT方法。

給定一個nn module的PyTorch圖形，我們首先將其轉(zhuǎn)換為Exported program。Exported program是在PyTorch 2.1中引入的一種新表示形式，它包含了Torch FX圖形和狀態(tài)字典兩部分。其中，Torch FX圖形包含了模型的張量計算，狀態(tài)字典用于存儲參數(shù)和緩沖區(qū)。這個轉(zhuǎn)換是通過使用Dynamo.trace API來完成的。此API是對Torch.export的封裝，并且除此之外，它還支持靜態(tài)和動態(tài)輸入。我們的追蹤器API還執(zhí)行一些附加的分解操作，以便將您的模型優(yōu)化轉(zhuǎn)換為TensorRT格式。

一旦我們獲得了Exported program，我們AOT方法的主要API就是Dynamo.compile。這個API將這些Exported program轉(zhuǎn)換為優(yōu)化的TensorRT圖形模塊。在幕后，這個API執(zhí)行l(wèi)owering、分區(qū)和轉(zhuǎn)換為TensorRT的操作，這些都是我們編譯器的核心組件，正如您在演示的開頭所看到的。一旦編譯完成，輸出的結(jié)果是一個包含TensorRT圖形模塊的TorchFX模塊。

下面的代碼段列出了該API的簡單用法。一旦您聲明了您的模型，只需將其傳遞給dynamo.trace，然后是dynamo.compile，該函數(shù)將返回優(yōu)化后的TensorRT圖模塊。

TensorRT期望圖中每個動態(tài)輸入都有一系列形狀。

在右邊的圖中，您可以看到我們可以使用torch_tensorrt.input API提供這些形狀范圍。與剛剛看到的jit流程的區(qū)別是您可以提供這個形狀范圍。這樣做的好處是，如果輸入形狀在提供的范圍內(nèi)發(fā)生更改，您無需重新編譯即可進行推理。靜態(tài)是序列化的主要好處之一。

為了總結(jié)我們到目前為止所見到的內(nèi)容，根據(jù)您的PyTorch圖形，我們使用我們的trace API生成導(dǎo)出的程序，然后使用Dynamo.compile API進行編譯。最后，您將獲得一個TorchFX圖模塊，其中包含TensorRT優(yōu)化的引擎?，F(xiàn)在，您可以使用Dynamo.Serialize API，將這些圖形模塊對象轉(zhuǎn)換為編程腳本或?qū)С龀绦虻谋硎拘问?，并隨后保存到磁盤上。同樣，右側(cè)的代碼片段非常易于使用。一旦您從Dynamo.compile中獲得了TensorRT模型，只需使用模型及其輸入調(diào)用serialize API即可。

以下是我們目前所見的內(nèi)容的概述。我們能夠處理復(fù)雜的Python代碼。通過重新編譯支持動態(tài)形狀，將已編譯的模塊序列化。此外，我們還能夠通過我們的輸入API支持動態(tài)形狀，并且將它們序列化為.script或?qū)С龅某绦?，但它無法處理任何圖形斷點。

然而，這兩者之間存在一些重要的相似之處。它們都經(jīng)歷類似的圖形優(yōu)化以進行高性能推斷。Torch TensorRT在PyTorch框架中以兩個關(guān)鍵路徑提供了優(yōu)化的推理方式。

結(jié)論和未來工作
Torch-TensorRT通過兩個關(guān)鍵路徑在PyTorch中提供了優(yōu)化的推理：

對于JIT工作流和復(fù)雜模型，采用基于torch.compile的方法

對于AoT工作流和序列化，采用強大的基于Exported program的方法

未來的工作：

在兩個路徑上提供對動態(tài)形狀的全面支持

改進對編譯Exported program的序列化支持

通過支持額外的精度和優(yōu)化，提高模型性能