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

10.1. Defining A Loop

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

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

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

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

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

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

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

一個(gè)循環(huán)可以有多個(gè)IIteratorLayer 、 IRecurrenceLayer和ILoopOutputLayer ，并且最多可以有兩個(gè)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)（請(qǐng)參閱嵌套循環(huán)）和其他條件構(gòu)造（請(qǐng)參閱條件嵌套）。

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

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

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

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

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

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

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

2 3 5

4 6 8

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

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

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

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

loop -》addLoopOutput（ t， LoopOutput：：kLAST_VALUE）輸出t的最后一個(gè)值，其中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形式都需要第二個(gè)輸入，這是一個(gè) 0D INT32 形狀張量，用于指定新輸出維度的長(zhǎng)度。當(dāng)長(zhǎng)度大于迭代次數(shù)時(shí)，額外的元素包含任意值。第二個(gè)輸入，例如u ，應(yīng)使用ILoopOutputLayer：：setInput（1， u ）設(shè)置。

最后，還有IRecurrenceLayer 。它的第一個(gè)輸入指定初始輸出值，第二個(gè)輸入指定下一個(gè)輸出值。第一個(gè)輸入必須來自循環(huán)外；第二個(gè)輸入通常來自循環(huán)內(nèi)部。例如，這個(gè) 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）;

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

10.2. Formal Semantics

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

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

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

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

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

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

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

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

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

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

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

10.3. Nested Loops

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

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

10.4. Limitations

引用多個(gè)動(dòng)態(tài)維度的循環(huán)可能會(huì)占用意外的內(nèi)存量。 在一個(gè)循環(huán)中，內(nèi)存的分配就像所有動(dòng)態(tài)維度都取任何這些維度的最大值一樣。例如，如果一個(gè)循環(huán)引用兩個(gè)維度為［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ò)。例如，請(qǐng)參閱sampleCharRNN方法SampleCharRNNLoop：：addLSTMCell 。循環(huán) API 讓您可以表達(dá)一般的循環(huán)網(wǎng)絡(luò)，而不是局限于IRNNLayer和IRNNv2Layer中的預(yù)制單元。

關(guān)于作者

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

審核編輯：郭婷