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

大綱

Torch TensorRT介紹

JIT編譯與AOT編譯方法

兩種方法的異同點(diǎn)

詳細(xì)要點(diǎn)

1. Torch TensorRT介紹

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

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

2. 兩種編譯方法

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

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

3. 編譯流程

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

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

4. JIT編譯細(xì)節(jié)

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

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

5. AOT編譯細(xì)節(jié)

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

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

支持序列化模型

6. 兩種方法的異同

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

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

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

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

我將在今天的用例中分享一些方法，包括Model performance和core implementation。這意味著，如果你有一個(gè)單一的模型，并且你用我們的兩種方法之一進(jìn)行優(yōu)化，你可以得到相似的性能和核心軟件實(shí)現(xiàn)?，F(xiàn)在你可能會(huì)問，為什么我要選擇其中之一？因?yàn)椴煌牟渴饒?chǎng)景需要不同的方法。

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

對(duì)于AoT工作流和序列化，采用強(qiáng)大的基于Exported program的方法

未來(lái)的工作：

在兩個(gè)路徑上提供對(duì)動(dòng)態(tài)形狀的全面支持

改進(jìn)對(duì)編譯Exported program的序列化支持

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