NVIDIA TensorRT 支持循環(huán)結(jié)構(gòu)，這對于循環(huán)網(wǎng)絡(luò)很有用。 TensorRT 循環(huán)支持掃描輸入張量、張量的循環(huán)定義以及“掃描輸出”和“最后一個值”輸出。

10.1. Defining A Loop

循環(huán)由循環(huán)邊界層（loop boundary layers）定義。

ITripLimitLayer指定循環(huán)迭代的次數(shù)。

IIteratorLayer使循環(huán)能夠迭代張量。

IRecurrenceLayer指定一個循環(huán)定義。

ILoopOutputLayer指定循環(huán)的輸出。

每個邊界層都繼承自類ILoopBoundaryLayer ，該類有一個方法getLoop（）用于獲取其關(guān)聯(lián)的ILoop 。 ILoop對象標識循環(huán)。具有相同ILoop的所有循環(huán)邊界層都屬于該循環(huán)。

下圖描繪了循環(huán)的結(jié)構(gòu)和邊界處的數(shù)據(jù)流。循環(huán)不變張量可以直接在循環(huán)內(nèi)部使用，例如 FooLayer 所示。

一個循環(huán)可以有多個IIteratorLayer 、 IRecurrenceLayer和ILoopOutputLayer ，并且最多可以有兩個ITripLimitLayer ，如后面所述。沒有ILoopOutputLayer的循環(huán)沒有輸出，并由 TensorRT 優(yōu)化。

NVIDIA TensorRT 支持矩陣中的流控制結(jié)構(gòu)層部分描述了可用于循環(huán)內(nèi)部的 TensorRT 層。

內(nèi)部層可以自由使用在循環(huán)內(nèi)部或外部定義的張量。內(nèi)部可以包含其他循環(huán)（請參閱嵌套循環(huán)）和其他條件構(gòu)造（請參閱條件嵌套）。

要定義循環(huán)，首先，使用INetworkDefinition ：：addLoop方法創(chuàng)建一個ILoop對象。然后添加邊界層和內(nèi)部層。本節(jié)的其余部分描述了邊界層的特征，使用loop表示INetworkDefinition ：：addLoop返回的ILoop* 。

ITripLimitLayer支持計數(shù)循環(huán)和 while 循環(huán)。

loop -》addTripLimit（ t ，TripLimit：：kCOUNT）創(chuàng)建一個ITripLimitLayer ，其輸入t是指定循環(huán)迭代次數(shù)的 0D INT32 張量。

loop -》addTripLimit（ t ，TripLimit：：kWHILE）創(chuàng)建一個ITripLimitLayer ，其輸入t是一個 0D Bool 張量，用于指定是否應(yīng)該進行迭代。通常t要么是IRecurrenceLayer的輸出，要么是基于所述輸出的計算。

一個循環(huán)最多可以有一種限制。

IIteratorLayer支持在任何軸上向前或向后迭代。

loop -》addIterator（ t ）添加一個IIteratorLayer ，它在張量t的軸 0 上進行迭代。例如，如果輸入是矩陣：

2 3 5

4 6 8

第一次迭代的一維張量{2， 3， 5}和第二次迭代的{4， 6， 8} 。超出張量范圍的迭代是無效的。

loop -》addIterator（ t ， axis ）類似，但該層在給定的軸上迭代。例如，如果 axis=1 并且輸入是矩陣，則每次迭代都會傳遞矩陣的一列。

loop -》addIterator（ t ， axis，reverse ）類似，但如果reverse =true ，則該層以相反的順序產(chǎn)生其輸出。

ILoopOutputLayer支持三種形式的循環(huán)輸出：

loop -》addLoopOutput（ t， LoopOutput：：kLAST_VALUE）輸出t的最后一個值，其中t必須是IRecurrenceLayer的輸出。

loop-》 addLoopOutput（ t ，LoopOutput：：kCONCATENATE， axis ）將每次迭代的輸入串聯(lián)輸出到t 。例如，如果輸入是一維張量，第一次迭代的值為{ a，b，c} ，第二次迭代的值為{d，e，f} ， axis =0 ，則輸出為矩陣：

a b c

d e f

如果axis =1 ，則輸出為：

a d

b e

c f

loop-》 addLoopOutput（ t ，LoopOutput：：kREVERSE， axis ）類似，但顛倒了順序。 kCONCATENATE和kREVERSE形式都需要第二個輸入，這是一個 0D INT32 形狀張量，用于指定新輸出維度的長度。當長度大于迭代次數(shù)時，額外的元素包含任意值。第二個輸入，例如u ，應(yīng)使用ILoopOutputLayer：：setInput（1， u ）設(shè)置。

最后，還有IRecurrenceLayer 。它的第一個輸入指定初始輸出值，第二個輸入指定下一個輸出值。第一個輸入必須來自循環(huán)外；第二個輸入通常來自循環(huán)內(nèi)部。例如，這個 C++ 片段的 TensorRT 模擬：

for （int32_t i = j; 。..; i += k） 。..

可以通過這些調(diào)用創(chuàng)建，其中j和k是ITensor* 。

ILoop* loop = n.addLoop（）;

IRecurrenceLayer* iRec = loop-》addRecurrence（j）;

ITensor* i = iRec-》getOutput（0）;

ITensor* iNext = addElementWise（*i， *k，

ElementWiseOperation：：kADD）-》getOutput（0）;

iRec-》setInput（1， *iNext）;

第二個輸入是TensorRT允許后沿的唯一情況。如果刪除了這些輸入，則剩余的網(wǎng)絡(luò)必須是非循環(huán)的。

10.2. Formal Semantics

TensorRT 具有應(yīng)用語義，這意味著除了引擎輸入和輸出之外沒有可見的副作用。因為沒有副作用，命令式語言中關(guān)于循環(huán)的直覺并不總是有效。本節(jié)定義了 TensorRT 循環(huán)結(jié)構(gòu)的形式語義。

形式語義基于張量的惰性序列（lazy sequences）。循環(huán)的每次迭代對應(yīng)于序列中的一個元素。循環(huán)內(nèi)張量X的序列表示為? X 0， X 1， X 2， 。.. ? 。序列的元素被懶惰地評估，意思是根據(jù)需要。

IIteratorLayer（X）的輸出是? X［0］， X［1］， X［2］， 。.. ?其中X［i］表示在IIteratorLayer指定的軸上的下標。

IRecurrenceLayer（X，Y）的輸出是? X， Y0， Y1， Y2， 。.. ? 。 的輸入和輸出取決于LoopOutput的類型。

kLAST_VALUE ：輸入是單個張量X ，對于 n-trip 循環(huán)，輸出是X n 。

kCONCATENATE ：第一個輸入是張量X，第二個輸入是標量形狀張量Y。結(jié)果是X0， X1， X2， 。.. Xn-1與后填充（如有必要）連接到Y(jié)指定的長度。如果Y 《 n則為運行時錯誤。 Y是構(gòu)建時間常數(shù)。注意與IIteratorLayer的反比關(guān)系。 IIteratorLayer將張量映射到一系列子張量；帶有kCONCATENATE的ILoopOutputLayer將一系列子張量映射到一個張量。

kREVERSE ：類似于kCONCATENATE ，但輸出方向相反。

ILoopOutputLayer的輸出定義中的n值由循環(huán)的ITripLimitLayer確定：

對于計數(shù)循環(huán)，它是迭代計數(shù)，表示ITripLimitLayer的輸入。

對于 while 循環(huán)，它是最小的 n 使得$X_n$為假，其中X是ITripLimitLayer的輸入張量的序列。

非循環(huán)層的輸出是層功能的順序應(yīng)用。例如，對于一個二輸入非循環(huán)層F（X，Y） = ? f（X 0 ， Y 0 ）， f（X 1 ， Y 1 ）， f（X 2 ， Y 2 ）。.. ? 。如果一個張量來自循環(huán)之外，即是循環(huán)不變的，那么它的序列是通過復(fù)制張量來創(chuàng)建的。

10.3. Nested Loops

TensorRT 從數(shù)據(jù)流中推斷出循環(huán)的嵌套。例如，如果循環(huán) B 使用在循環(huán) A 中定義的值，則 B 被認為嵌套在 A 中。

TensorRT 拒絕循環(huán)沒有干凈嵌套的網(wǎng)絡(luò)，例如如果循環(huán) A 使用循環(huán) B 內(nèi)部定義的值，反之亦然。

10.4. Limitations

引用多個動態(tài)維度的循環(huán)可能會占用意外的內(nèi)存量。 在一個循環(huán)中，內(nèi)存的分配就像所有動態(tài)維度都取任何這些維度的最大值一樣。例如，如果一個循環(huán)引用兩個維度為［4，x，y］和［6，y］的張量，則這些張量的內(nèi)存分配就像它們的維度是［4，max（x，y），max（x ，y）］和［6，max（x，y）］ 。

帶有kLAST_VALUE的LoopOutputLayer的輸入必須是IRecurrenceLayer的輸出。 循環(huán) API 僅支持 FP32 和 FP16 精度。

10.5. Replacing IRNNv2Layer With Loops

IRNNv2Layer在 TensorRT 7.2.1 中已棄用，并將在 TensorRT 9.0 中刪除。使用循環(huán) API 合成循環(huán)子網(wǎng)絡(luò)。例如，請參閱sampleCharRNN方法SampleCharRNNLoop：：addLSTMCell 。循環(huán) API 讓您可以表達一般的循環(huán)網(wǎng)絡(luò)，而不是局限于IRNNLayer和IRNNv2Layer中的預(yù)制單元。

關(guān)于作者

Ken He 是 NVIDIA 企業(yè)級開發(fā)者社區(qū)經(jīng)理 & 高級講師，擁有多年的 GPU 和人工智能開發(fā)經(jīng)驗。自 2017 年加入 NVIDIA 開發(fā)者社區(qū)以來，完成過上百場培訓(xùn)，幫助上萬個開發(fā)者了解人工智能和 GPU 編程開發(fā)。在計算機視覺，高性能計算領(lǐng)域完成過多個獨立項目。并且，在機器人和無人機領(lǐng)域，有過豐富的研發(fā)經(jīng)驗。對于圖像識別，目標的檢測與跟蹤完成過多種解決方案。曾經(jīng)參與 GPU 版氣象模式GRAPES，是其主要研發(fā)者。

審核編輯：郭婷