編寫 TorchScript 轉換器¶

背景¶

在 JIT IR 中，操作表示為圖中的節點。節點有輸入和輸出，由 torch::jit::Values 表示，它們是流入和流出節點的資料的型別化抽象表示。TensorRT 透過使用 nvinfer1::ILayers 和 nvinfer1::ITensors 來表示其圖，它們分別對應於節點和值。轉換器的目標是建立新的 ILayers 和子圖，以執行節點指定的操作，並將生成的 ITensors 和 Values 關聯起來。

轉換器¶

轉換器應該是函式，它們將使用輸入列表（可以是 nvinfer1::ITensors 或 torch::jit::IValues）來構建與 LibTorch 操作等效的層。

轉換器可以使用 RegisterNodeConversionPatterns 輔助類進行註冊，您例項化一個 RegisterNodeConversionPatterns 物件，並呼叫其上的 pattern 函式（如下所示），該函式接受一個字串，描述將觸發轉換器執行的操作的函式 schema，以及一個執行實際轉換的 lambda 或函式。

請注意，pattern 函式可以鏈式呼叫。

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。
輸入將按照函式 schema 的順序提供。

轉換器的職責¶

必須保證 args 是可以直接解包 Arg 聯合型別的特定型別，無需檢查；通常，輸入張量引數可以預期為 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()），您將使用這些 API 將層新增到 TRT 層，並記錄節點輸出與靜態值或 TensorRT 層輸出的配對。

引數¶

提供給轉換器的引數是 nvinfer1::ITensors 和 torch::jit::IValues 的可檢查聯合型別（即 TensorRT 圖中的抽象資料流和靜態值）。對於節點的每個輸入值，您都可以保證會獲得相應的引數。它們按照函式 schema 的順序提供。可以預期輸入（即在 PyTorch 中傳遞給模組 forward 函式的引數）將是 ITensors，但如果您不確定，Arg 類也有機制在解包前安全地檢查引數。Arg 類還具有深度解包方法，如果您確定安全，可以直接獲取 IValue 中的底層資料。如果 IValue 可能為 None，您還可以傳入一個備用值。IValues 已被擴充套件，僅在 TensorLists 的情況下能夠包含圍繞 ITensors 的包裝器。您可以使用類似於以下模式從 IValue 獲取 ITensor：ivalue.toCustomClass<TensorContainer>()->tensor()。您可以使用 ivalue.isTensor() 或 ivalue.isCustomClass() 判斷 IValue 是否包含 Tensor 或 ITensor。

權重¶

權重在構建時使用，因此必須保證所有權重在轉換階段結束前始終存活。TensorRT 也有自己的權重結構來儲存權重。轉換器可以使用一個圍繞此類的包裝器，該包裝器抽象了許多細節。

權重包裝類可以接受 at::Tensors 或單個值（目前）。在構建這些權重時，您還需要傳入轉換上下文，因為在內部，權重類將分配由轉換上下文管理的記憶體來儲存張量資料的副本。當轉換上下文的解構函式被銷燬時，這些資料就會被釋放，因此轉換器無需真正考慮它。

輸入資料的形狀會生成元資料，這在與 TensorRT 互動時非常有用，例如輸入對映數量、輸出對映數量和核心形狀。

其他建議¶

在處理權重和其他靜態值時，您可以受益於完整的 aten 庫。這意味著您可以在轉換期間進行大量工作以實現高效轉換。一個很好的例子是 batch_norm 轉換器，它在建立 TensorRT 層之前與 PyTorch 一起執行操作融合。