在 C++ 中使用 Torch-TensorRT¶
如果您還沒有這樣做，請按照安裝說明獲取庫的 tar 包。
在 C++ 中使用 Torch-TensorRT¶
Torch-TensorRT C++ API 接受 TorchScript 模組（由 torch.jit.script 或 torch.jit.trace 生成）作為輸入，並返回一個 Torchscript 模組（使用 TensorRT 最佳化）。但是，C++ API 將不支援 Dynamo 編譯工作流，但支援對 FX 和 Dynamo 工作流執行由 torch.jit.trace 生成的已編譯 FX GraphModules。
請參閱在 Python 中建立 TorchScript 模組一節來生成 torchscript 圖。
[Torch-TensorRT 快速入門] 使用 torchtrtc 編譯 TorchScript 模組¶
開始使用 Torch-TensorRT 並檢查您的模型是否無需額外工作即可受支援的一種簡單方法是透過 torchtrtc 執行它。 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 儲存我們的跟蹤（traced）或指令碼化（scripted）模組，它會將 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.com.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 的運算子庫非常龐大，因此有些運算子編譯器本身並不支援。您可以採用上面所示的組合技術，將完全受 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）運算子（aten::flatten）的 Torch-TensorRT 構建版本編譯圖，您可能會看到以下錯誤：
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() 的例項，它會在全域性轉換器登錄檔中註冊轉換器，將一個函式 schema（例如 aten::flatten.using_ints(Tensor self, int start_dim=0, int end_dim=-1) -> (Tensor)）與一個 lambda 函式關聯起來。這個 lambda 函式將接收轉換狀態、要轉換的節點/操作以及節點的所有輸入，其副作用是在 TensorRT 網路中生成一個新的層。引數以 TensorRT ITensors 和 Torch IValues 可檢查聯合的向量形式傳遞，順序與 schema 中列出的引數順序一致。
下面是一個 aten::flatten 轉換器的實現示例，我們可以在應用程式中使用它。在轉換器實現中，您可以完全訪問 Torch 和 TensorRT 庫。因此，例如，我們可以透過在 PyTorch 中直接執行操作來快速獲取輸出大小，而不是像我們下面為這個展平轉換器那樣自己實現完整的計算。
#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 擴充套件模板將您的轉換器庫包裝成一個 .so 檔案，您可以在 Python 應用程式中使用 ctypes.CDLL() 載入它。
您可以在貢獻者文件（writing_converters）中找到關於編寫轉換器所有細節的更多資訊。如果您有很多轉換器實現，請考慮將它們貢獻回上游，我們歡迎拉取請求（PR），這將極大地造福社群。
