• 文档 >
  • 部署 Torch-TensorRT 程序
快捷方式

部署 Torch-TensorRT 程序

在编译和保存 Torch-TensorRT 程序后,不再严格依赖于完整的 Torch-TensorRT 库。运行已编译程序所需的一切只是运行时库。因此,除了将完整的 Torch-TensorRT 编译器与您的应用程序一起发布外,还有几种部署程序的选项。

Torch-TensorRT 包 / libtorchtrt.so

程序编译后,您可以使用标准的 PyTorch API 来运行它。唯一的要求是在 Python 中导入该包或在 C++ 中链接它。

运行时库

C++ 发行版中附带了 libtorchtrt_runtime.so。该库仅包含运行 Torch-TensorRT 程序所必需的组件。您可以链接 libtorchtrt_runtime.so 到您的部署程序中,或使用 DL_OPENLD_PRELOAD,而无需链接 libtorchtrt.so 或导入 torch_tensorrt。对于 Python,您可以使用 torch.ops.load_library("libtorchtrt_runtime.so") 加载运行时。然后,您可以像往常一样通过 PyTorch API 继续使用程序。

注意

如果您要链接 libtorchtrt_runtime.so,使用以下标志可能会有帮助:-Wl,--no-as-needed -ltorchtrt -Wl,--as-needed,因为对于大多数 Torch-TensorRT 运行时应用程序,不存在对 Torch-TensorRT 运行时中任何内容的直接符号依赖。

关于如何使用 libtorchtrt_runtime.so 的示例可以在这里找到:https://github.com/pytorch/TensorRT/tree/master/examples/torchtrt_aoti_example

插件库

如果您将 Torch-TensorRT 用作 TensorRT 引擎的转换器,并且您的引擎使用了 Torch-TensorRT 提供的插件,Torch-TensorRT 会附带一个库 libtorchtrt_plugins.so,其中包含了 Torch-TensorRT 在编译期间使用的 TensorRT 插件的实现。该库可以通过 DL_OPENLD_PRELOAD 加载,与其他 TensorRT 插件库类似。

多设备安全模式

多设备安全模式是 Torch-TensorRT 中的一项设置,它允许用户决定运行时是否在每次推理调用之前检查设备一致性。

启用多设备安全模式时,每次推理调用都会产生不可忽略的固定成本,因此现在默认是禁用的。它可以通过以下既可作为上下文管理器又可作为便捷函数的工具来控制。

# Enables Multi Device Safe Mode
torch_tensorrt.runtime.set_multi_device_safe_mode(True)

# Disables Multi Device Safe Mode [Default Behavior]
torch_tensorrt.runtime.set_multi_device_safe_mode(False)

# Enables Multi Device Safe Mode, then resets the safe mode to its prior setting
with torch_tensorrt.runtime.set_multi_device_safe_mode(True):
    ...

TensorRT 要求每个引擎都与调用它的活动线程中的 CUDA 上下文相关联。因此,如果活动线程中的设备发生变化(例如在同一 Python 进程中调用多个 GPU 上的引擎时),安全模式将使 Torch-TensorRT 显示警报并相应地切换 GPU。如果未启用安全模式,引擎设备和 CUDA 上下文设备可能会不匹配,从而导致程序崩溃。

一种在不同 GPU 上管理多个 TRT 引擎,同时又不因多设备安全模式而牺牲性能的技术是使用 Python 线程。每个线程负责单个 GPU 上的所有 TRT 引擎,并且每个线程上的默认 CUDA 设备对应于其负责的 GPU(可通过 torch.cuda.set_device(...) 设置)。通过这种方式,可以在同一个 Python 脚本中使用多个线程,而无需切换 CUDA 上下文并产生性能开销。

Cudagraphs 模式

Cudagraphs 模式是 Torch-TensorRT 中的一项设置,它允许用户决定运行时是否在某些情况下使用 cudagraphs 来加速推理。

Cudagraphs 可以通过减少内核开销来加速某些模型,具体如[此处](https://pytorch.ac.cn/blog/accelerating-pytorch-with-cuda-graphs/)所述。

# Enables Cudagraphs Mode
torch_tensorrt.runtime.set_cudagraphs_mode(True)

# Disables Cudagraphs Mode [Default Behavior]
torch_tensorrt.runtime.set_cudagraphs_mode(False)

# Enables Cudagraphs Mode, then resets the mode to its prior setting
with torch_tensorrt.runtime.enable_cudagraphs(trt_module):
    ...

在当前的实现中,使用新的输入形状(例如在动态形状的情况下)将导致 cudagraph 重新录制。Cudagraph 录制通常不会对延迟产生很大影响,未来的改进包括为多种输入形状缓存 cudagraph。

动态输出分配模式

动态输出分配是 Torch-TensorRT 中的一项功能,它允许动态分配 TensorRT 引擎的输出缓冲区。这对于具有动态输出形状的模型非常有用,尤其是那些具有数据依赖形状的操作。动态输出分配模式不能与 CUDA Graphs 或预分配输出功能一起使用。在没有动态输出分配的情况下,输出缓冲区是根据输入大小推断出的输出形状来分配的。

在两种情况下会启用动态输出分配:

1. 在编译时已识别出模型至少有一个 TensorRT 子图需要动态输出分配。这些模型将自动启用运行时模式(并带有日志记录),并且与 CUDA Graphs 等其他运行时模式不兼容。

转换器可以通过 requires_output_allocator=True 声明它们生成的子图将需要输出分配器,从而强制任何使用该转换器的模型自动使用输出分配器运行时模式。例如,

@dynamo_tensorrt_converter(
    torch.ops.aten.nonzero.default,
    supports_dynamic_shapes=True,
    requires_output_allocator=True,
)
def aten_ops_nonzero(
    ctx: ConversionContext,
    target: Target,
    args: Tuple[Argument, ...],
    kwargs: Dict[str, Argument],
    name: str,
) -> Union[TRTTensor, Sequence[TRTTensor]]:
    ...
  1. 用户可以通过 torch_tensorrt.runtime.enable_output_allocator 上下文管理器手动启用动态输出分配模式。

# Enables Dynamic Output Allocation Mode, then resets the mode to its prior setting
with torch_tensorrt.runtime.enable_output_allocator(trt_module):
    ...

不使用 Python 部署 Torch-TensorRT 程序

AOT-Inductor

AOTInductor 是 TorchInductor 的一个专门版本,旨在处理导出的 PyTorch 模型,对其进行优化,并生成共享库以及其他相关工件。这些编译后的工件专门为非 Python 环境中的部署而设计,这些环境通常用于服务器端的推理部署。

Torch-TensorRT 能够像在 Python 中一样,加速 AOTInductor 导出中的子图。

dynamo_model = torch_tensorrt.compile(model, ir="dynamo", arg_inputs=[...])
torch_tensorrt.save(
    dynamo_model,
    file_path=os.path.join(os.getcwd(), "model.pt2"),
    output_format="aot_inductor",
    retrace=True,
    arg_inputs=[...],
)

然后,这个工件可以在 C++ 应用程序中加载并执行,而无需 Python 依赖。

#include <iostream>
#include <vector>

#include "torch/torch.h"
#include "torch/csrc/inductor/aoti_package/model_package_loader.h"

int main(int argc, const char* argv[]) {
// Check for correct number of command-line arguments
std::string trt_aoti_module_path = "model.pt2";

if (argc == 2) {
    trt_aoti_module_path = argv[1];
}

    std::cout << trt_aoti_module_path << std::endl;

    // Get the path to the TRT AOTI model package from the command line
    c10::InferenceMode mode;

    torch::inductor::AOTIModelPackageLoader loader(trt_aoti_module_path);
    // Assume running on CUDA
    std::vector<torch::Tensor> inputs = {torch::randn({8, 10}, at::kCUDA)};
    std::vector<torch::Tensor> outputs = loader.run(inputs);
    std::cout << "Result from the first inference:"<< std::endl;
    std::cout << outputs << std::endl;

    // The second inference uses a different batch size and it works because we
    // specified that dimension as dynamic when compiling model.pt2.
    std::cout << "Result from the second inference:"<< std::endl;
    // Assume running on CUDA
    std::cout << loader.run({torch::randn({1, 10}, at::kCUDA)}) << std::endl;

    return 0;
}

注意:与 Python 类似,在运行时,不会使用 Torch-TensorRT API 来操作模型。因此,通常需要额外的标志来确保 libtorchtrt_runtime.so 不会被优化掉(见上文)。

请参阅://examples/torchtrt_aoti_example 以获取此工作流程的完整端到端演示

TorchScript

TorchScript 是 PyTorch 的一个旧版编译器堆栈,它包含一个用于 TorchScript 程序的无 Python 解释器。历史上,Torch-TensorRT 使用它来在没有 Python 的情况下执行模型。即使在过渡到 TorchDynamo 之后,TorchScript 解释器仍然可以继续用于在 Python 之外运行带有 TensorRT 引擎的 PyTorch 模型。

dynamo_model = torch_tensorrt.compile(model, ir="dynamo", arg_inputs=[...])
ts_model = torch.jit.trace(dynamo_model, inputs=[...])
torch.jit.save(ts_model, os.path.join(os.getcwd(), "model.ts"),)

然后,这个工件可以在 C++ 应用程序中加载并执行,而无需 Python 依赖。

#include <fstream>
#include <iostream>
#include <memory>
#include <sstream>
#include <vector>
#include "torch/script.h"

int main(int argc, const char* argv[]) {
    if (argc < 2) {
        std::cerr << "usage: samplertapp <path-to-pre-built-trt-ts module>\n";
        return -1;
    }

    std::string trt_ts_module_path = argv[1];

    torch::jit::Module trt_ts_mod;
    try {
        // Deserialize the ScriptModule from a file using torch::jit::load().
        trt_ts_mod = torch::jit::load(trt_ts_module_path);
    } catch (const c10::Error& e) {
        std::cerr << "error loading the model from : " << trt_ts_module_path << std::endl;
        return -1;
    }

    std::cout << "Running TRT engine" << std::endl;
    std::vector<torch::jit::IValue> trt_inputs_ivalues;
    trt_inputs_ivalues.push_back(at::randint(-5, 5, {1, 3, 5, 5}, {at::kCUDA}).to(torch::kFloat32));
    torch::jit::IValue trt_results_ivalues = trt_ts_mod.forward(trt_inputs_ivalues);
    std::cout << "==================TRT outputs================" << std::endl;
    std::cout << trt_results_ivalues << std::endl;
    std::cout << "=============================================" << std::endl;
    std::cout << "TRT engine execution completed. " << std::endl;
}

注意:与 Python 类似,在运行时,不会使用 Torch-TensorRT API 来操作模型。因此,通常需要额外的标志来确保 libtorchtrt_runtime.so 不会被优化掉(见上文)。

请参阅://examples/torchtrt_runtime_example 以获取此工作流程的完整端到端演示

文档

访问全面的 PyTorch 开发者文档

查看文档

教程

为初学者和高级开发者提供深入的教程

查看教程

资源

查找开发资源并让您的问题得到解答

查看资源