在 C++ 中使用 Torch-TensorRT

如果您尚未取得程式庫的壓縮檔，請依照安裝中的說明進行操作

在 C++ 中使用 Torch-TensorRT

Torch-TensorRT C++ API 接受 TorchScript 模組（從 torch.jit.script 或 torch.jit.trace 生成）作為輸入，並傳回 Torchscript 模組（使用 TensorRT 最佳化）。 C++ API 將不支援 Dynamo 編譯工作流程，但是，支援執行 torch.jit.trace 編譯的 FX GraphModules 用於 FX 和 Dyanmo 工作流程。

請參閱 在 Python 中建立 TorchScript 模組 章節，以產生 torchscript 圖。

[Torch-TensorRT 快速入門] 使用 torchtrtc 編譯 TorchScript 模組

要開始使用 Torch-TensorRT 並檢查您的模型是否無需額外工作即可支援的簡單方法，是透過 torchtrtc 執行它。它幾乎支援命令列中的所有編譯器功能，包括訓練後量化（給定先前建立的校正快取）。例如，我們可以透過設定我們偏好的運算精度和輸入大小來編譯我們的 lenet 模型。這個新的 TorchScript 檔案可以載入到 Python 中（注意：在載入這些編譯後的模組之前，您需要 import torch_tensorrt，因為編譯器擴充了 PyTorch 的反序列化器和執行時環境來執行編譯後的模組）。

❯ torchtrtc -p f16 lenet_scripted.ts trt_lenet_scripted.ts "(1,1,32,32)"

❯ python3
Python 3.6.9 (default, Apr 18 2020, 01:56:04)
[GCC 8.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> import torch_tensorrt
>>> ts_model = torch.jit.load(“trt_lenet_scripted.ts”)
>>> ts_model(torch.randn((1,1,32,32)).to(“cuda”).half())

您可以在這裡了解更多關於 torchtrtc 的用法：torchtrtc

在 C++ 中使用 TorchScript

如果我們正在開發一個要使用 C++ 部署的應用程式，我們可以使用 torch.jit.save 來儲存我們追蹤或編寫的模組，這會將 TorchScript 程式碼、權重和其他資訊序列化到一個套件中。這也是我們對 Python 的依賴結束的地方。

torch_script_module.save("lenet.jit.pt")

從這裡我們現在可以在 C++ 中載入我們的 TorchScript 模組

#include <torch/script.h> // One-stop header.

#include <iostream>
#include <memory>

int main(int argc, const char* argv[]) {
    torch::jit::Module module;
    try {
        // Deserialize the ScriptModule from a file using torch::jit::load().
        module = torch::jit::load("<PATH TO SAVED TS MOD>");
    }
    catch (const c10::Error& e) {
        std::cerr << "error loading the model\n";
        return -1;
    }

    std::cout << "ok\n";

如果您願意，您可以使用 PyTorch / LibTorch 在 C++ 中進行完整的訓練和推論，您甚至可以在 C++ 中定義您的模組，並存取支援 PyTorch 的相同強大的張量庫。（更多資訊：https://pytorch.dev.org.tw/cppdocs/）。例如，我們可以像這樣使用我們的 LeNet 模組進行推論

mod.eval();
torch::Tensor in = torch::randn({1, 1, 32, 32});
auto out = mod.forward(in);

並且在 GPU 上運行

mod.eval();
mod.to(torch::kCUDA);
torch::Tensor in = torch::randn({1, 1, 32, 32}, torch::kCUDA);
auto out = mod.forward(in);

正如您所看到的，它與 Python API 非常相似。當您呼叫 forward 方法時，您會調用 PyTorch JIT 編譯器，它將優化並運行您的 TorchScript 程式碼。

在 C++ 中使用 Torch-TensorRT 編譯

我們也到了可以使用 Torch-TensorRT 編譯和優化我們的模組的階段，但不是以 JIT 方式，而是必須提前 (AOT) 執行，即在我們開始進行實際推論工作之前，因為優化模組需要一些時間，每次運行模組甚至第一次運行它都沒有意義。

載入我們的模組後，我們可以將其饋送到 Torch-TensorRT 編譯器中。在這樣做時，我們必須提供一些關於預期輸入大小的資訊，並且配置任何其他設定。

#include "torch/script.h"
#include "torch_tensorrt/torch_tensorrt.h"
...

    mod.to(at::kCUDA);
    mod.eval();
    std::vector<torch_tensorrt::core::ir::Input> inputs{torch_tensorrt::core::ir::Input({1, 3, 224, 224})};
    torch_tensorrt::ts::CompileSpec cfg(inputs);
    auto trt_mod = torch_tensorrt::ts::compile(mod, cfg);
    auto in = torch::randn({1, 3, 224, 224}, {torch::kCUDA});
    auto out = trt_mod.forward({in});

就是這樣！現在圖主要不是使用 JIT 編譯器，而是使用 TensorRT 運行（儘管我們使用 JIT 執行時環境執行該圖）。

我們還可以設定諸如運算精度之類的設定以在 FP16 中運行。

#include "torch/script.h"
#include "torch_tensorrt/torch_tensorrt.h"
...

    mod.to(at::kCUDA);
    mod.eval();

    auto in = torch::randn({1, 3, 224, 224}, {torch::kCUDA}).to(torch::kHALF);
    std::vector<torch_tensorrt::core::ir::Input> inputs{torch_tensorrt::core::ir::Input({1, 3, 224, 224})};
    torch_tensorrt::ts::CompileSpec cfg(inputs);
    cfg.enable_precisions.insert(torch::kHALF);
    auto trt_mod = torch_tensorrt::ts::compile(mod, cfg);
    auto out = trt_mod.forward({in});

現在我們正在 FP16 精度下運行該模組。然後，您可以儲存該模組以供以後載入。

trt_mod.save("<PATH TO SAVED TRT/TS MOD>")

Torch-TensorRT 編譯的 TorchScript 模組的載入方式與正常的 TorchScript 模組相同。請確保您的部署應用程式已連結到 libtorchtrt.so

#include "torch/script.h"
#include "torch_tensorrt/torch_tensorrt.h"

int main(int argc, const char* argv[]) {
    torch::jit::Module module;
    try {
        // Deserialize the ScriptModule from a file using torch::jit::load().
        module = torch::jit::load("<PATH TO SAVED TRT/TS MOD>");
    }
    catch (const c10::Error& e) {
        std::cerr << "error loading the model\n";
        return -1;
    }

    torch::Tensor in = torch::randn({1, 1, 32, 32}, torch::kCUDA);
    auto out = mod.forward(in);

    std::cout << "ok\n";
}

如果您想要儲存 Torch-TensorRT 產生的引擎以在 TensorRT 應用程式中使用，您可以使用 ConvertGraphToTRTEngine API。

#include "torch/script.h"
#include "torch_tensorrt/torch_tensorrt.h"
...

    mod.to(at::kCUDA);
    mod.eval();

    auto in = torch::randn({1, 3, 224, 224}, {torch::kCUDA}).to(torch::kHALF);

    std::vector<torch_tensorrt::core::ir::Input> inputs{torch_tensorrt::core::ir::Input({1, 3, 224, 224})};
    torch_tensorrt::ts::CompileSpec cfg(inputs);
    cfg.enabled_precisions.insert(torch::kHALF);
    auto trt_mod = torch_tensorrt::ts::convert_method_to_trt_engine(mod, "forward", cfg);
    std::ofstream out("/tmp/engine_converted_from_jit.trt");
    out << engine;
    out.close();

幕後

當模組被提供給 Torch-TensorRT 時，編譯器首先將您在上面看到的圖映射到這樣的圖

graph(%input.2 : Tensor):
    %2 : Float(84, 10) = prim::Constant[value=<Tensor>]()
    %3 : Float(120, 84) = prim::Constant[value=<Tensor>]()
    %4 : Float(576, 120) = prim::Constant[value=<Tensor>]()
    %5 : int = prim::Constant[value=-1]() # x.py:25:0
    %6 : int[] = prim::Constant[value=annotate(List[int], [])]()
    %7 : int[] = prim::Constant[value=[2, 2]]()
    %8 : int[] = prim::Constant[value=[0, 0]]()
    %9 : int[] = prim::Constant[value=[1, 1]]()
    %10 : bool = prim::Constant[value=1]() # ~/.local/lib/python3.6/site-packages/torch/nn/modules/conv.py:346:0
    %11 : int = prim::Constant[value=1]() # ~/.local/lib/python3.6/site-packages/torch/nn/functional.py:539:0
    %12 : bool = prim::Constant[value=0]() # ~/.local/lib/python3.6/site-packages/torch/nn/functional.py:539:0
    %self.classifier.fc3.bias : Float(10) = prim::Constant[value= 0.0464  0.0383  0.0678  0.0932  0.1045 -0.0805 -0.0435 -0.0818  0.0208 -0.0358 [ CUDAFloatType{10} ]]()
    %self.classifier.fc2.bias : Float(84) = prim::Constant[value=<Tensor>]()
    %self.classifier.fc1.bias : Float(120) = prim::Constant[value=<Tensor>]()
    %self.feat.conv2.weight : Float(16, 6, 3, 3) = prim::Constant[value=<Tensor>]()
    %self.feat.conv2.bias : Float(16) = prim::Constant[value=<Tensor>]()
    %self.feat.conv1.weight : Float(6, 1, 3, 3) = prim::Constant[value=<Tensor>]()
    %self.feat.conv1.bias : Float(6) = prim::Constant[value= 0.0530 -0.1691  0.2802  0.1502  0.1056 -0.1549 [ CUDAFloatType{6} ]]()
    %input0.4 : Tensor = aten::_convolution(%input.2, %self.feat.conv1.weight, %self.feat.conv1.bias, %9, %8, %9, %12, %8, %11, %12, %12, %10) # ~/.local/lib/python3.6/site-packages/torch/nn/modules/conv.py:346:0
    %input0.5 : Tensor = aten::relu(%input0.4) # ~/.local/lib/python3.6/site-packages/torch/nn/functional.py:1063:0
    %input1.2 : Tensor = aten::max_pool2d(%input0.5, %7, %6, %8, %9, %12) # ~/.local/lib/python3.6/site-packages/torch/nn/functional.py:539:0
    %input0.6 : Tensor = aten::_convolution(%input1.2, %self.feat.conv2.weight, %self.feat.conv2.bias, %9, %8, %9, %12, %8, %11, %12, %12, %10) # ~/.local/lib/python3.6/site-packages/torch/nn/modules/conv.py:346:0
    %input2.1 : Tensor = aten::relu(%input0.6) # ~/.local/lib/python3.6/site-packages/torch/nn/functional.py:1063:0
    %x.1 : Tensor = aten::max_pool2d(%input2.1, %7, %6, %8, %9, %12) # ~/.local/lib/python3.6/site-packages/torch/nn/functional.py:539:0
    %input.1 : Tensor = aten::flatten(%x.1, %11, %5) # x.py:25:0
    %27 : Tensor = aten::matmul(%input.1, %4)
    %28 : Tensor = trt::const(%self.classifier.fc1.bias)
    %29 : Tensor = aten::add_(%28, %27, %11)
    %input0.2 : Tensor = aten::relu(%29) # ~/.local/lib/python3.6/site-packages/torch/nn/functional.py:1063:0
    %31 : Tensor = aten::matmul(%input0.2, %3)
    %32 : Tensor = trt::const(%self.classifier.fc2.bias)
    %33 : Tensor = aten::add_(%32, %31, %11)
    %input1.1 : Tensor = aten::relu(%33) # ~/.local/lib/python3.6/site-packages/torch/nn/functional.py:1063:0
    %35 : Tensor = aten::matmul(%input1.1, %2)
    %36 : Tensor = trt::const(%self.classifier.fc3.bias)
    %37 : Tensor = aten::add_(%36, %35, %11)
    return (%37)
(CompileGraph)

現在，該圖已經從一個模組集合轉換而來，每個模組都管理自己的參數，變成一個單一圖，其參數被內聯到該圖中，並且所有操作都被佈局。 Torch-TensorRT 也執行了許多優化和映射，以使該圖更容易轉換為 TensorRT。從這裡開始，編譯器可以通過跟隨圖中的資料流來組裝 TensorRT 引擎。

當圖構造階段完成時，Torch-TensorRT 會產生一個序列化的 TensorRT 引擎。從這裡開始，根據 API，此引擎將返回給使用者或進入圖構造階段。在這裡，Torch-TensorRT 創建一個 JIT 模組來執行 TensorRT 引擎，該引擎將由 Torch-TensorRT 執行時環境實例化和管理。

這是編譯完成後您獲得的回溯圖

graph(%self_1 : __torch__.lenet, %input_0 : Tensor):
    %1 : ...trt.Engine = prim::GetAttr[name="lenet"](%self_1)
    %3 : Tensor[] = prim::ListConstruct(%input_0)
    %4 : Tensor[] = trt::execute_engine(%3, %1)
    %5 : Tensor = prim::ListUnpack(%4)
    return (%5)

您可以看到執行引擎的呼叫，在提取包含引擎的屬性並構建輸入列表後，然後將張量返回給使用者。

使用不受支援的運算子

Torch-TensorRT 是一個新的函式庫，而 PyTorch 運算子函式庫非常龐大，因此會有編譯器本身不支援的運算 (op)。您可以使用上面顯示的組合技術來製作完全支援 Torch-TensorRT 的模組，以及不支援的模組，並在部署應用程式中將這些模組縫合在一起，或者您可以為缺少的運算子註冊轉換器。

您可以使用 torch_tensorrt::CheckMethodOperatorSupport(const torch::jit::Module& module, std::string method_name) api 檢查支援，而無需經過完整的編譯管道，以查看哪些運算子不受支援。torchtrtc 會在開始編譯之前使用此方法自動檢查模組，並印出不受支援的運算子列表。

註冊自定義轉換器

運算透過模組化轉換器映射到 TensorRT，模組化轉換器是一種函式，它從 JIT 圖中取得一個節點，並在 TensorRT 中產生等效的層或子圖。 Torch-TensorRT 附帶一個儲存在註冊表中的這些轉換器的函式庫，它將根據正在解析的節點執行。例如，aten::relu(%input0.4) 指令將觸發 relu 轉換器在其上運行，從而在 TensorRT 圖中產生一個激活層。但由於此函式庫並不詳盡，您可能需要編寫自己的程式碼才能使 Torch-TensorRT 支援您的模組。

與 Torch-TensorRT 發行版一起提供的是內部核心 API 標頭。因此，您可以訪問轉換器註冊表並為您需要的運算添加一個轉換器。

例如，如果我們嘗試使用不支援 flatten 運算的 Torch-TensorRT 版本（aten::flatten）編譯圖，您可能會看到此錯誤

terminate called after throwing an instance of 'torch_tensorrt::Error'
what():  [enforce fail at core/conversion/conversion.cpp:109] Expected converter to be true but got false
Unable to convert node: %input.1 : Tensor = aten::flatten(%x.1, %11, %5) # x.py:25:0 (conversion.AddLayer)
Schema: aten::flatten.using_ints(Tensor self, int start_dim=0, int end_dim=-1) -> (Tensor)
Converter for aten::flatten requested, but no such converter was found.
If you need a converter for this operator, you can try implementing one yourself
or request a converter: https://www.github.com/NVIDIA/Torch-TensorRT/issues

我們可以在應用程式中註冊此運算子的轉換器。建構轉換器所需的所有工具都可以透過包含 torch_tensorrt/core/conversion/converters/converters.h 來匯入。 我們首先建立一個自我註冊類別 torch_tensorrt::core::conversion::converters::RegisterNodeConversionPatterns() 的實例，它會在全域轉換器註冊表中註冊轉換器，將一個函數模式（例如 aten::flatten.using_ints(Tensor self, int start_dim=0, int end_dim=-1) -> (Tensor)）與一個 lambda 關聯起來。該 lambda 函數會接收轉換的狀態、要轉換的節點/運算，以及節點的所有輸入，並產生 TensorRT 網路中的一個新層作為副作用。引數會以 TensorRT ITensors 和 Torch IValues 的可檢查聯集向量形式傳遞，其順序與模式中列出的引數順序相同。

以下是在我們的應用程式中可以使用的 aten::flatten 轉換器的實作範例。您可以在轉換器實作中完全存取 Torch 和 TensorRT 函式庫。 因此，舉例來說，我們可以透過在 PyTorch 中執行運算來快速取得輸出大小，而無需像我們為此 flatten 轉換器所做的那樣，自己實作完整的計算過程。

#include "torch/script.h"
#include "torch_tensorrt/torch_tensorrt.h"
#include "torch_tensorrt/core/conversion/converters/converters.h"

static auto flatten_converter = torch_tensorrt::core::conversion::converters::RegisterNodeConversionPatterns()
    .pattern({
        "aten::flatten.using_ints(Tensor self, int start_dim=0, int end_dim=-1) -> (Tensor)",
        [](torch_tensorrt::core::conversion::ConversionCtx* ctx,
           const torch::jit::Node* n,
           torch_tensorrt::core::conversion::converters::args& args) -> bool {
            auto in = args[0].ITensor();
            auto start_dim = args[1].unwrapToInt();
            auto end_dim = args[2].unwrapToInt();
            auto in_shape = torch_tensorrt::core::util::toVec(in->getDimensions());
            auto out_shape = torch::flatten(torch::rand(in_shape), start_dim, end_dim).sizes();

            auto shuffle = ctx->net->addShuffle(*in);
            shuffle->setReshapeDimensions(torch_tensorrt::core::util::toDims(out_shape));
            shuffle->setName(torch_tensorrt::core::util::node_info(n).c_str());

            auto out_tensor = ctx->AssociateValueAndTensor(n->outputs()[0], shuffle->getOutput(0));
            return true;
        }
    });

int main() {
    ...

若要在 Python 中使用此轉換器，建議使用 PyTorch 的C++ / CUDA Extension 範本，將您的轉換器函式庫封裝成一個 .so 檔案，您可以使用 Python 應用程式中的 ctypes.CDLL() 來載入它。

您可以在貢獻者文件中找到關於編寫轉換器的所有細節的更多資訊（writing_converters）。 如果您發現自己擁有大量的轉換器實作函式庫，請考慮將它們向上游推送，我們歡迎 PR，並且這對社群來說也是一件好事。