撰寫轉換器¶

背景¶

在 JIT IR 中，運算以圖形中的節點表示。節點具有輸入和輸出，由 torch::jit::Values 表示，這些值是流入和流出節點的數據的類型化抽象表示。TensorRT 通過使用 nvinfer1::ILayers 和 nvinfer1::ITensors 來表示其圖形，這兩者分別類似於節點和值。轉換器的目標是創建新的 ILayers 和子圖，這些 ILayers 和子圖執行節點指定的運算，並將生成的 ITensors 和 Values 關聯在一起。

轉換器¶

轉換器應該是函數，這些函數將使用輸入列表（nvinfer1::ITensors 或 torch::jit::IValues）來構造等效於 LibTorch 運算的層。

可以使用 RegisterNodeConversionPatterns 輔助類別來註冊轉換器，在其中實例化一個 RegisterNodeConversionPatterns 物件並在其上呼叫模式函數（如下所示），該函數採用一個字串，該字串描述將導致轉換器運行的運算的函數架構，以及一個將執行實際轉換的 lambda 或函數

請注意，模式函數可以鏈接

auto acthardtanh TORCHTRT_UNUSED = RegisterNodeConversionPatterns()
    .pattern({
        "aten::hardtanh(Tensor self, Scalar min_val=-1, Scalar max_val=1) -> (Tensor)",
        [](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
            auto in = args[0].ITensor();
            auto min = args[1].unwrapToDouble();
            auto max = args[2].unwrapToDouble();

            auto new_layer = ctx->net->addActivation(*in, nvinfer1::ActivationType::kCLIP);
            TORCHTRT_CHECK(new_layer, "Unable to create layer for aten::hardtanh");

            new_layer->setAlpha(min);
            new_layer->setBeta(max);

            new_layer->setName(util::node_info(n).c_str());
            auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], new_layer->getOutput(0));

            LOG_DEBUG("Output shape: " << out_tensor->getDimensions());
            return true;
        }
    });

轉換器合約¶

轉換器的保證¶

在 args 中，將為每個節點輸入值都有一個條目，即 ITensor 或 IValue
輸入將根據函數架構按順序提供

轉換器的責任¶

必須保證 Args 是可以在不檢查的情況下解開 Arg union 的類型，通常輸入張量參數可以預計是 ITensors
必須保證任何權重或靜態值在轉換時間結束之前都有效
1. 一個有用的工具是下面描述的 Weights 輔助類別
預計轉換器會為節點的每個輸出生成一個 IValue 或 ITensor。編譯器將檢查這一點，如果存在沒有關聯 ITensors 或 IValues 的 Values，則會產生警告。
輸出必須註釋
1. 在轉換上下文的 value_tensor_map 中，JIT 節點的輸出值與新的 TRT 層輸出張量之間必須有關聯
命名您的層
1. 當我們可以追蹤哪些層和節點彼此對應時，除錯起來就容易多了。我們目前使用的系統是使用節點的“節點資訊”作為層的名稱
命名您的張量
1. 使用輸出值的除錯名稱作為新 ITensor 的名稱（同樣是為了除錯）

轉換上下文¶

轉換上下文維護轉換的狀態，它管理網路定義、兩個映射（一個存儲 Values 和 IValues 之間的關聯（evaluated_value_map），另一個存儲 Values 和 ITensors 之間的關聯）以及任何需要存活到轉換結束的記憶體。您將在轉換器中與之交互的主要 API 是直接訪問網路定義以添加層 ctx->net 和數據關聯函數 ctx->AssociateValueAndTensor() 和 ctx->AssociateValueAndIValue()，您將使用它們來添加層到 TRT 層並記錄節點輸出和靜態值或 TensorRT 層輸出的對。

參數¶

提供給轉換器的參數是 nvinfer1::ITensors 和 torch::jit::IValues 的可檢查 union（即 TensorRT 圖形中的抽象數據流和靜態值）。您可以保證您將為節點的每個輸入值都有一些參數。它們按照函數架構的順序提供。可以預期輸入（意指將在 PyTorch 中傳遞給模組的 forward 函數的參數）將是 ITensors，但如果您不確定，Arg 類別還具有在解開之前安全檢查參數的機制。Args 還具有深度解開方法，如果您知道安全的話，可以直接獲取 IValue 中的底層數據。如果 IValue 有可能是 None，您也可以傳遞一個後備值。IValues 已被擴展為能夠保存 ITensors 的包裝器，僅在 TensorLists 的情況下。您可以通過類似於以下模式的方式從 IValue 中獲取 ITensor：ivalue.toCustomClass<TensorContainer>()->tensor()。您可以通過使用 ivalue.isTensor() 或 ivalue.isCustomClass() 來判斷 IValue 是否包含 Tensor 或 ITensor。

權重¶

權重在構建時使用，因此需要保證任何權重都存活到轉換階段結束。TensorRT 還使用自己的權重結構來保存權重。轉換器可以使用圍繞此類別的包裝器，它抽象了很多東西。

權重包裝器類別可以接受 at::Tensors 或單個值（目前）。在構建這些權重時，您還需要傳遞轉換上下文，因為在內部，權重類別將分配由轉換上下文管理的記憶體來存儲張量數據的副本。當轉換上下文解構函數被銷毀時，這些數據將被釋放，因此轉換器實際上不需要考慮它。

元數據是根據輸入數據的形狀生成的，這在與 TensorRT 交互時變得很有用，例如輸入映射的數量、輸出映射的數量和核心形狀。

其他建議¶

在處理權重和其他靜態值時，您可以利用完整的 aten 函數庫。這意味著您可以在轉換期間做很多工作來提高轉換效率。一個很好的例子是 batch_norm 轉換器，其中轉換器在創建 TensorRT 層之前使用 PyTorch 進行運算融合。