使用循环#

NVIDIA TensorRT 支持类似循环的结构，这对于循环神经网络非常有用。TensorRT 循环支持扫描输入张量、张量的循环定义，以及扫描输出和最后值输出。

定义循环#

循环边界层定义了一个循环。

  • ITripLimitLayer 指定循环迭代的次数。

  • IIteratorLayer 使循环能够迭代张量。

  • IRecurrenceLayer 指定循环定义。

  • ILoopOutputLayer 指定循环的输出。

每个边界层都继承自类 ILoopBoundaryLayer，该类具有方法 getLoop()，用于获取其关联的 ILoop。ILoop 对象标识循环，并且具有相同 ILoop 的所有循环边界层都属于该循环。

下图描述了循环结构和边界处的数据流。循环不变张量可以直接在循环内部使用，如 FooLayer 所示。

A TensorRT loop is set by loop boundary layers. Dataflow can leave the loop only by ``ILoopOutputLayer``. The only back edges allowed are the second input to ``IRecurrenceLayer``.

稍后将解释，一个循环可以有多个 IIteratorLayer、IRecurrenceLayer 和 ILoopOutputLayer，以及最多两个 ITripLimitLayers。没有 ILoopOutputLayer 的循环没有输出，并由 TensorRT 优化。

内部层可以自由使用在循环内部或外部定义的张量。内部可以包含其他循环（请参阅 嵌套循环）和其他条件结构（请参阅 嵌套和循环）。

要定义循环，首先使用 INetworkDefinition method::addLoop 创建 ILoop 对象。然后，添加边界层和内部层。本节的其余部分描述了边界层的功能，使用循环来表示 INetworkDefinition::addLoop 返回的 ILoop*。

ITripLimitLayer 支持计数循环和 while 循环。

  • loop->addTripLimit(t,TripLimit::kCOUNT) 创建一个 ITripLimitLayer，其输入 t 是一个 0D INT32 张量，用于指定循环迭代的次数。

  • loop->addTripLimit(t,TripLimit::kWHILE) 创建一个 ITripLimitLayer，其输入 t 是一个 0D Bool 张量，用于指定是否应进行迭代。通常，t 是 IRecurrenceLayer 的输出或基于所述输出的计算。

一个循环最多可以包含每种限制类型中的一个。

IIteratorLayer 支持正向或反向迭代任何轴。

  • loop->addIterator(t) 添加一个 IIteratorLayer，该层迭代张量 t 的轴 0。例如，如果输入是矩阵

    [[2,3,5],
 [4,6,8]]

则在第一次迭代时，输出是 1D 张量 {2, 3, 5}，在第二次迭代时，输出是 {4, 6, 8}。迭代超出张量的边界是无效的。

  • loop->addIterator(t,axis) 类似，但该层迭代给定的轴。例如，如果 axis=1 并且输入是矩阵，则每次迭代都会传递一个矩阵列。

  • loop->addIterator(t,axis,reverse) 类似，但如果 reverse=true，则该层以相反的顺序生成其输出。

ILoopOutputLayer 支持三种形式的循环输出

  • loop->addLoopOutput(t,LoopOutput::kLAST_VALUE) 输出 t 的最后一个值，其中 t 必须是 IRecurrenceLayer 的输出。

  • loop->addLoopOutput(t,LoopOutput::kCONCATENATE,axis) 输出每次迭代的 t 输入的串联。例如，如果输入是 1D 张量，第一次迭代的值为 {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) 类似，但反转了顺序。

ILoopOutputLayer 的 kCONCATENATE 和 kREVERSE 形式都需要第二个输入，即 0D INT32 形状张量，用于指定新输出维度的长度。当长度超过迭代次数时，额外的元素包含任意值。例如，第二个输入 u 应使用 ILoopOutputLayer::setInput(1,u) 进行设置。

最后，还有 IRecurrenceLayer。它的第一个输入指定初始输出值，第二个输入指定下一个输出值。第一个输入必须来自循环外部；第二个输入通常来自内部。例如，TensorRT 类似于以下 C++ 代码片段

for (int32_t i = j; ...; i += k) ...

这些调用可以创建如下内容，其中 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);

IRecurrenceLayer 的第二个输入是 TensorRT 允许反向边的唯一情况。如果删除此类输入，则剩余的网络必须是非循环的。

形式语义#

TensorRT 具有应用语义，这意味着除了引擎输入和输出之外，没有可见的副作用。由于没有副作用，因此从命令式语言中获得的关于循环的直觉并不总是有效。本节定义了 TensorRT 循环结构的形式语义。

形式语义基于张量的惰性序列。循环的每次迭代都对应于序列中的一个元素。循环内张量 X 的序列表示为 \(\left\langle X_{0},X_{1},X_{2}, ... \right\rangle\)。序列的元素是惰性求值的，这意味着按需求值。

IIteratorLayer(X) 的输出是 \(\left\langle X\left[ 0 \right],X\left[ 1 \right],X\left[ 2 \right], ... \right\rangle\)，其中 X[i] 表示在为 IIteratorLayer 指定的轴上的下标。

IRecurrenceLayer(X,Y) 的输出是 \(\left\langle \text{X},Y_{0},Y_{1},Y_{2} ... \right\rangle\)。

ILoopOutputLayer 的输入和输出取决于 LoopOutput 的类型。

  • kLAST_VALUE：输入是单个张量 X，输出是 n 次循环的 \(X_{n}\)。

  • kCONCATENATE：第一个输入是张量 X，第二个输入是标量形状张量 Y。结果是 \(X_{0},X_{1},X_{2}, ... X_{n-1}\) 的串联，并在必要时进行后填充，以达到 Y 指定的长度。如果 Y < n，则会发生运行时错误。Y 是构建时常量。请注意与 IIteratorLayer 的反向关系。IIteratorLayer 将张量映射到子张量序列；ILoopOutputLayer 与 kCONCATENATE 将子张量序列映射到张量。

  • kREVERSE：类似于 kCONCATENATE，但输出方向相反。

ILoopOutputLayer 输出的定义中 n 的值由循环的 ITripLimitLayer 确定

  • 对于计数循环，它是迭代计数，即 ITripLimitLayer 的输入。

  • 对于 while 循环，它是最小的 n，使得 \(X_{n}\) 为假，其中 X 是 ITripLimitLayer 输入张量的序列。

非循环层的输出是该层函数的序列式应用。例如，对于双输入非循环层 \(F\left( X,Y \right)=\left\langle f\left( X_{0},Y_{0} \right),f\left( X_{1},Y_{1} \right),f\left( X_{2},Y_{2} \right) \right\rangle\)。如果张量来自循环外部，即循环不变，则通过复制张量来创建它的序列。

嵌套循环#

TensorRT 从数据流中推断循环的嵌套。例如，如果循环 B 使用在循环 A内部定义的值，则 B 被认为嵌套在 A 内部。

TensorRT 拒绝循环未干净嵌套的网络，例如，如果循环 A 使用循环 B 内部定义的值，反之亦然。

局限性#

引用多个动态维度的循环可能会占用超出预期的内存量。在循环中，内存分配就好像所有动态维度都采用这些维度中的最大值一样。例如，如果一个循环引用了维度为 [4,x,y] 和 [6,y] 的两个张量，则这些张量的内存分配就好像它们的维度为 [4,max(x,y),max(x,y)] 和 [6,max(x,y)]。

带有 kLAST_VALUE 的 LoopOutputLayer 的输入必须来自 IRecurrenceLayer 的输出。

循环 API 仅支持 FP32 和 FP16 精度。

用循环替换 IRNNv2Layer#

IRNNv2Layer 在 TensorRT 7.2.1 中已弃用，并在 TensorRT 10.0 中移除。使用循环 API 合成循环子网络。例如，请参阅 sampleCharRNN，方法 SampleCharRNNLoop::addLSTMCell。使用循环 API，您可以表达通用的循环网络，而不是局限于 IRNNLayer 和 IRNNv2Layer 中的预制单元。

在本页中