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

11.1. Defining A Conditional

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

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

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

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

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

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

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

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

IIfConditional* simpleIf = network->addIfConditional();

IIfConditional ::setCondition方法接受一個(gè)參數(shù)：條件張量。這個(gè) 0D 布爾張量（標(biāo)量）可以由網(wǎng)絡(luò)中的早期層動(dòng)態(tài)計(jì)算。它用于決定執(zhí)行哪個(gè)分支。IConditionLayer有一個(gè)輸入（條件）并且沒(méi)有輸出，因?yàn)樗蓷l件實(shí)現(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 不支持實(shí)現(xiàn)條件分支的子圖抽象，而是使用IIfConditionalInputLayer和IIfConditionalOutputLayer來(lái)定義條件的邊界。

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

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

• IIfConditionalOutputLayer抽象了if條件的單個(gè)輸出。它有兩個(gè)輸入：來(lái)自真子圖的輸出（輸入索引 0）和來(lái)自假子圖的輸出（輸入索引 1）。IIfConditionalOutputLayer的輸出可以被認(rèn)為是最終輸出的占位符，最終輸出將在運(yùn)行時(shí)確定。IIfConditionalOutputLayer的作用類似于傳統(tǒng) SSA 控制流圖中的 $Φ(Phi)$ 函數(shù)節(jié)點(diǎn)。它的語(yǔ)義是：選擇真子圖或假子圖的輸出。IIfConditional的所有輸出都必須源自IIfConditionalOutputLayer實(shí)例。沒(méi)有輸出的 if 條件對(duì)網(wǎng)絡(luò)的其余部分沒(méi)有影響，因此，它被認(rèn)為是病態(tài)的。兩個(gè)分支（子圖）中的每一個(gè)也必須至少有一個(gè)輸出。if-conditional的輸出可以標(biāo)記為網(wǎng)絡(luò)的輸出，除非if-conditional嵌套在另一個(gè)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ò)評(píng)估策略，其中僅在需要分支輸出的值時(shí)才執(zhí)行分支層（屬于條件子圖的層）。在條件執(zhí)行中，無(wú)論是真分支還是假分支都被執(zhí)行并允許改變網(wǎng)絡(luò)狀態(tài)。

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

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

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

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

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

在圖表 C 中，輸入層 I1 被完全移除，這將 T3 移到條件之外。 T2 的輸入被重新配置以創(chuàng)建合法網(wǎng)絡(luò)，并且 T2 也移出真實(shí)分支。當(dāng)條件評(píng)估為true時(shí)，條件不計(jì)算任何內(nèi)容，因?yàn)檩敵鲆呀?jīng)被急切地計(jì)算（但它確實(shí)將條件相關(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)。

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

例如，請(qǐng)參閱條件示例以了解如何指定嵌套。

11.4. Limitations

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

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

11.5. Conditional Examples

11.5.1. Simple If-Conditional

下面的例子展示了如何實(shí)現(xiàn)一個(gè)簡(jiǎn)單的條件，它有條件地對(duì)兩個(gè)張量執(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")