TensorRT條件用于實現(xiàn)網(wǎng)絡(luò)子圖的條件執(zhí)行

11.1. Defining A Conditional

if-conditional 由條件邊界層定義：

• IConditionLayer表示predicate 并指定條件是應(yīng)該執(zhí)行真分支（then-branch）還是假分支（else-branch）。
• IIfConditionalInputLayer指定兩個條件分支之一的輸入。
• IIfConditionalOutputLayer指定條件的輸出。

每個邊界層都繼承自IIfConditionalBoundaryLayer類，該類具有獲取其關(guān)聯(lián)IIfConditional的方法getConditional()。IIfConditional實例標(biāo)識條件。所有具有相同IIfConditional的條件邊界層都屬于該條件。

條件必須恰好有一個IConditionLayer實例、零個或多個IIfConditionalInputLayer實例，以及至少一個IIfConditionalOutputLayer實例。

IIfConditional實現(xiàn)了一個if-then-else流控制結(jié)構(gòu)，該結(jié)構(gòu)提供基于動態(tài)布爾輸入的網(wǎng)絡(luò)子圖的條件執(zhí)行。它由一個布爾標(biāo)量predicate condition和兩個分支子圖定義：一個trueSubgraph在condition評估為true時執(zhí)行，一個falseSubgraph在condition評估為false時執(zhí)行

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

真分支和假分支都必須定義，類似于許多編程語言中的三元運(yùn)算符。

要定義if-conditional，使用方法INetworkDefinition::addIfConditional創(chuàng)建一個IIfConditional實例，然后添加邊界層和分支層。

IIfConditional* simpleIf = network->addIfConditional();

IIfConditional ::setCondition方法接受一個參數(shù)：條件張量。這個 0D 布爾張量（標(biāo)量）可以由網(wǎng)絡(luò)中的早期層動態(tài)計算。它用于決定執(zhí)行哪個分支。IConditionLayer有一個輸入（條件）并且沒有輸出，因為它由條件實現(xiàn)在內(nèi)部使用。

// 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 不支持實現(xiàn)條件分支的子圖抽象，而是使用IIfConditionalInputLayer和IIfConditionalOutputLayer來定義條件的邊界。

• IIfConditionalInputLayer將單個輸入抽象為IIfConditional的一個或兩個分支子圖。特定IIfConditionalInputLayer的輸出可以同時提供兩個分支。then-branch和else-branch的輸入不需要是相同的類型和形狀，每個分支可以獨立地包含零個或多個輸入。IIfConditionalInputLayer是可選的，用于控制哪些層將成為分支的一部分（請參閱條件執(zhí)行）。如果分支的所有輸出都不依賴于IIfConditionalInputLayer實例，則該分支為空。當(dāng)條件為false時沒有要評估的層時，空的else-branch可能很有用，并且網(wǎng)絡(luò)評估應(yīng)按照條件進(jìn)行（請參閱條件示例）。

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

• IIfConditionalOutputLayer抽象了if條件的單個輸出。它有兩個輸入：來自真子圖的輸出（輸入索引 0）和來自假子圖的輸出（輸入索引 1）。IIfConditionalOutputLayer的輸出可以被認(rèn)為是最終輸出的占位符，最終輸出將在運(yùn)行時確定。IIfConditionalOutputLayer的作用類似于傳統(tǒng) SSA 控制流圖中的 $Φ(Phi)$ 函數(shù)節(jié)點。它的語義是：選擇真子圖或假子圖的輸出。IIfConditional的所有輸出都必須源自IIfConditionalOutputLayer實例。沒有輸出的 if 條件對網(wǎng)絡(luò)的其余部分沒有影響，因此，它被認(rèn)為是病態(tài)的。兩個分支（子圖）中的每一個也必須至少有一個輸出。if-conditional的輸出可以標(biāo)記為網(wǎng)絡(luò)的輸出，除非if-conditional嵌套在另一個if-conditional或循環(huán)中。

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

下圖提供了 if 條件抽象模型的圖形表示。綠色矩形表示條件的內(nèi)部，僅限于NVIDIA TensorRT 支持矩陣中的LayersFor Flow-Control Constructs部分中列出的層類型。

11.2. Conditional Execution

網(wǎng)絡(luò)層的條件執(zhí)行是一種網(wǎng)絡(luò)評估策略，其中僅在需要分支輸出的值時才執(zhí)行分支層（屬于條件子圖的層）。在條件執(zhí)行中，無論是真分支還是假分支都被執(zhí)行并允許改變網(wǎng)絡(luò)狀態(tài)。

相反，在斷定執(zhí)行中，真分支和假分支都被執(zhí)行，并且只允許其中之一改變網(wǎng)絡(luò)評估狀態(tài)，具體取決于條件斷定的值（即僅其中一個的輸出）子圖被饋送到以下層。

條件執(zhí)行有時稱為惰性求值，斷定執(zhí)行有時稱為急切求值。IIfConditionalInputLayer的實例可用于指定急切調(diào)用哪些層以及延遲調(diào)用哪些層。這是通過從每個條件輸出開始向后跟蹤網(wǎng)絡(luò)層來完成的。依賴于至少一個IIfConditionalInputLayer輸出的數(shù)據(jù)層被認(rèn)為是條件內(nèi)部的，因此被延遲評估。在沒有IIfConditionalInputLayer實例添加到條件條件的極端情況下，所有層都被急切地執(zhí)行，類似于ISelectLayer。

下面的三個圖表描述了IIfConditionalInputLayer放置的選擇如何控制執(zhí)行調(diào)度。

在圖 A 中，真分支由 3 層（T1、T2、T3）組成。當(dāng)條件評估為true時，這些層會延遲執(zhí)行。

在圖 B 中，輸入層 I1 放置在層 T1 之后，它將 T1 移出真實分支。在評估 if 結(jié)構(gòu)之前，T1 層急切地執(zhí)行。

在圖表 C 中，輸入層 I1 被完全移除，這將 T3 移到條件之外。 T2 的輸入被重新配置以創(chuàng)建合法網(wǎng)絡(luò)，并且 T2 也移出真實分支。當(dāng)條件評估為true時，條件不計算任何內(nèi)容，因為輸出已經(jīng)被急切地計算（但它確實將條件相關(guān)輸入復(fù)制到其輸出）。

11.3. Nesting and Loops

條件分支可以嵌套其他條件，也可以嵌套循環(huán)。循環(huán)可以嵌套條件。與循環(huán)嵌套一樣，TensorRT 從數(shù)據(jù)流中推斷條件和循環(huán)的嵌套。例如，如果條件 B 使用在循環(huán) A 內(nèi)定義的值，則 B 被認(rèn)為嵌套在 A 內(nèi)。

真分支中的層與假分支中的層之間不能有交叉邊，反之亦然。換句話說，一個分支的輸出不能依賴于另一個分支中的層。

例如，請參閱條件示例以了解如何指定嵌套。

11.4. Limitations

兩個真/假子圖分支中的輸出張量數(shù)必須相同。來自分支的每個輸出張量的類型和形狀必須相同。

請注意，這比 ONNX 規(guī)范更受限制，ONNX 規(guī)范要求真/假子圖具有相同數(shù)量的輸出并使用相同的輸出數(shù)據(jù)類型，但允許不同的輸出形狀。

11.5. Conditional Examples

11.5.1. Simple If-Conditional

下面的例子展示了如何實現(xiàn)一個簡單的條件，它有條件地對兩個張量執(zhí)行算術(shù)運(yùn)算。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);

11.5.2. Exporting from PyTorch

以下示例展示了如何將腳本化的 PyTorch 代碼導(dǎo)出到 ONNX。函數(shù)sum_even中的代碼執(zhí)行嵌套在循環(huán)中的 if 條件。

import torch.onnx
import torch
import tensorrt as trt
import numpy as np

TRT_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 += c
    return s

class 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")