TensorRT之条件控制

点击此处加入NVIDIA开发者计划

NVIDIA TensorRT 支持条件 if-then-else 流控制。 TensorRT 条件用于实现网络子图的条件执行。

1. Defining A Conditional

if-conditional 由条件边界层定义：

IConditionLayer表示predicate 并指定条件是应该执行真分支（then-branch）还是假分支（else-branch）。
IIfConditionalInputLayer指定两个条件分支之一的输入。
IIfConditionalOutputLayer指定条件的输出。

每个边界层都继承自IIfConditionalBoundaryLayer类，该类具有获取其关联IIfConditional的方法getConditional() 。 IIfConditional实例标识条件。所有具有相同IIfConditional的条件边界层都属于该条件。

条件必须恰好有一个IConditionLayer实例、零个或多个IIfConditionalInputLayer实例，以及至少一个IIfConditionalOutputLayer实例。

IIfConditional实现了一个 if-then-else 流控制结构，该结构提供基于动态布尔输入的网络子图的条件执行。它由一个布尔标量predicate condition和两个分支子图定义：一个trueSubgraph在condition评估为true时执行，一个falseSubgraph在condition评估为false时执行

If condition is true then: output = trueSubgraph(trueInputs);
Elseoutput = falseSubgraph(falseInputs);
Emit output

真分支和假分支都必须定义，类似于许多编程语言中的三元运算符。

要定义if-conditional，使用方法INetworkDefinition::addIfConditional创建一个IIfConditional实例，然后添加边界层和分支层。

IIfConditional* simpleIf = network->addIfConditional();

IIfConditional ::setCondition方法接受一个参数：条件张量。这个 0D 布尔张量（标量）可以由网络中的早期层动态计算。它用于决定执行哪个分支。 IConditionLayer有一个输入（条件）并且没有输出，因为它由条件实现在内部使用。

// Create a condition predicate that is also a network input.
auto cond = network->addInput("cond", DataType::kBOOL, Dims{0});
IConditionLayer* condition = simpleIf->setCondition(*cond);

TensorRT 不支持实现条件分支的子图抽象，而是使用IIfConditionalInputLayer和IIfConditionalOutputLayer来定义条件的边界。

IIfConditionalInputLayer将单个输入抽象为IIfConditional的一个或两个分支子图。特定IIfConditionalInputLayer的输出可以同时提供两个分支。then-branch 和 else-branch 的输入不需要是相同的类型和形状，每个分支可以独立地包含零个或多个输入。IIfConditionalInputLayer是可选的，用于控制哪些层将成为分支的一部分（请参阅条件执行）。如果分支的所有输出都不依赖于IIfConditionalInputLayer实例，则该分支为空。当条件为false时没有要评估的层时，空的 else-branch 可能很有用，并且网络评估应按照条件进行（请参阅条件示例）。

// Create an if-conditional input.
// x is some arbitrary Network tensor.
IIfConditionalInputLayer* inputX = simpleIf->addInput(*x);

IIfConditionalOutputLayer抽象了 if 条件的单个输出。它有两个输入：来自真子图的输出（输入索引 0）和来自假子图的输出（输入索引 1）。 IIfConditionalOutputLayer的输出可以被认为是最终输出的占位符，最终输出将在运行时确定。IIfConditionalOutputLayer的作用类似于传统 SSA 控制流图中的 Φ ( P h i ) Φ(Phi) Φ(Phi) 函数节点。它的语义是：选择真子图或假子图的输出。IIfConditional的所有输出都必须源自IIfConditionalOutputLayer实例。没有输出的 if 条件对网络的其余部分没有影响，因此，它被认为是病态的。两个分支（子图）中的每一个也必须至少有一个输出。 if-conditional 的输出可以标记为网络的输出，除非 if-conditional 嵌套在另一个 if-conditional 或循环中。

// trueSubgraph and falseSubgraph represent network subgraphs
IIfConditionalOutputLayer* outputLayer = simpleIf->addOutput(*trueSubgraph->getOutput(0), *falseSubgraph->getOutput(0));

下图提供了 if 条件抽象模型的图形表示。绿色矩形表示条件的内部，仅限于NVIDIA TensorRT 支持矩阵中的Layers For Flow-Control Constructs部分中列出的层类型。

[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-OVk6LMJz-1655821838052)(conditional_abstract_model.png)]

2. Conditional Execution

网络层的条件执行是一种网络评估策略，其中仅在需要分支输出的值时才执行分支层（属于条件子图的层）。在条件执行中，无论是真分支还是假分支都被执行并允许改变网络状态。

相反，在断定执行中，真分支和假分支都被执行，并且只允许其中之一改变网络评估状态，具体取决于条件断定的值（即仅其中一个的输出）子图被馈送到以下层。

条件执行有时称为惰性求值，断定执行有时称为急切求值。
IIfConditionalInputLayer的实例可用于指定急切调用哪些层以及延迟调用哪些层。这是通过从每个条件输出开始向后跟踪网络层来完成的。依赖于至少一个IIfConditionalInputLayer输出的数据层被认为是条件内部的，因此被延迟评估。在没有IIfConditionalInputLayer实例添加到条件条件的极端情况下，所有层都被急切地执行，类似于ISelectLayer 。

下面的三个图表描述了IIfConditionalInputLayer放置的选择如何控制执行调度。

[外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传(img-2DP5zqm2-1655821838052)(conditional_lazy_eager_execution.png)]

在图 A 中，真分支由 3 层（T1、T2、T3）组成。当条件评估为true时，这些层会延迟执行。

在图 B 中，输入层 I1 放置在层 T1 之后，它将 T1 移出真实分支。在评估 if 结构之前，T1 层急切地执行。

在图表 C 中，输入层 I1 被完全移除，这将 T3 移到条件之外。 T2 的输入被重新配置以创建合法网络，并且 T2 也移出真实分支。当条件评估为true时，条件不计算任何内容，因为输出已经被急切地计算（但它确实将条件相关输入复制到其输出）。

3. Nesting and Loops

条件分支可以嵌套其他条件，也可以嵌套循环。循环可以嵌套条件。与循环嵌套一样，TensorRT 从数据流中推断条件和循环的嵌套。例如，如果条件 B 使用在循环 A 内定义的值，则 B 被认为嵌套在 A 内。

真分支中的层与假分支中的层之间不能有交叉边，反之亦然。换句话说，一个分支的输出不能依赖于另一个分支中的层。

例如，请参阅条件示例以了解如何指定嵌套。

4. Limitations

两个真/假子图分支中的输出张量数必须相同。来自分支的每个输出张量的类型和形状必须相同。

请注意，这比 ONNX 规范更受限制，ONNX 规范要求真/假子图具有相同数量的输出并使用相同的输出数据类型，但允许不同的输出形状。

5. Conditional Examples

5.1. Simple If-Conditional

下面的例子展示了如何实现一个简单的条件，它有条件地对两个张量执行算术运算。
Conditional

condition = true
If condition is true:output = x + y
Else:output = x - y

Example

ITensor* addCondition(INetworkDefinition& n, bool predicate)
{// The condition value is a constant int32 input that is cast to boolean because TensorRT doesn't support boolean constant layers.static const Dims scalarDims = Dims{0, {}};static float constexpr zero{0};static float constexpr one{1};float* const val = predicate ? &one : &zero;ITensor* cond = n.addConstant(scalarDims, DataType::kINT32, val, 1})->getOutput(0);auto* cast = n.addIdentity(cond);cast->setOutputType(0, DataType::kBOOL);cast->getOutput(0)->setType(DataType::kBOOL);return cast->getOutput(0);
}IBuilder* builder = createInferBuilder(gLogger);
INetworkDefinition& n = *builder->createNetworkV2(0U);
auto x = n.addInput("x", DataType::kFLOAT, Dims{1, {5}});
auto y = n.addInput("y", DataType::kFLOAT, Dims{1, {5}});
ITensor* cond = addCondition(n, true);auto* simpleIf = n.addIfConditional();
simpleIf->setCondition(*cond);// Add input layers to demarcate entry into true/false branches.
x = simpleIf->addInput(*x)->getOutput(0);
y = simpleIf->addInput(*y)->getOutput(0);auto* trueSubgraph = n.addElementWise(*x, *y, ElementWiseOperation::kSUM)->getOutput(0);
auto* falseSubgraph = n.addElementWise(*x, *y, ElementWiseOperation::kSUB)->getOutput(0);auto* output = simpleIf->addOutput(*trueSubgraph, *falseSubgraph)->getOutput(0);
n.markOutput(*output);

5.2. Exporting from PyTorch

以下示例展示了如何将脚本化的 PyTorch 代码导出到 ONNX。函数sum_even中的代码执行嵌套在循环中的 if 条件。

import torch.onnx
import torch
import tensorrt as trt
import numpy as npTRT_LOGGER = trt.Logger(trt.Logger.WARNING)
EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)@torch.jit.script
def sum_even(items):s = torch.zeros(1, dtype=torch.float)for c in items:if c % 2 == 0:s += creturn sclass ExampleModel(torch.nn.Module):def __init__(self):super().__init__()def forward(self, items):return sum_even(items)def build_engine(model_file):builder = trt.Builder(TRT_LOGGER)network = builder.create_network(EXPLICIT_BATCH)config = builder.create_builder_config()parser = trt.OnnxParser(network, TRT_LOGGER)with open(model_file, 'rb') as model:assert parser.parse(model.read())return builder.build_engine(network, config)def export_to_onnx():items = torch.zeros(4, dtype=torch.float)example = ExampleModel()torch.onnx.export(example, (items), "example.onnx", verbose=False, opset_version=13, enable_onnx_checker=False, do_constant_folding=True)export_to_onnx()
build_engine("example.onnx")