TensorRT

项目介绍：Torch-TensorRT

什么是 Torch-TensorRT？

Torch-TensorRT 是一个强大的工具，能够将 TensorRT 的强大功能带入 PyTorch。它可以在 NVIDIA 平台上让任何 PyTorch 模型实现最佳推理性能。通过简单的一行代码，与常规的即时执行方式相比，它能够将推理延迟加速高达五倍。

安装指南

Torch-TensorRT 的稳定版本可以通过 PyPI 获得，只需运行以下命令：

pip install torch-tensorrt

而每晚构建的版本则可以从 PyTorch 包索引下载：

pip install --pre torch-tensorrt --index-url https://download.pytorch.org/whl/nightly/cu124

此外，Torch-TensorRT 也在 NVIDIA NGC PyTorch 容器中提供，这个容器包含了所有必要的依赖，并附带示例笔记本。

快速入门

选项 1：使用 torch.compile

可以在任何需要使用 torch.compile 的地方应用 Torch-TensorRT：

import torch
import torch_tensorrt

model = MyModel().eval().cuda() # 在这里定义您的模型
x = torch.randn((1, 3, 224, 224)).cuda() # 定义输入格式

optimized_model = torch.compile(model, backend="tensorrt")
optimized_model(x) # 第一次运行时编译

optimized_model(x) # 之后将变得非常迅速！

选项 2：导出和部署

如果想在提前优化模型并在 C++ 环境中部署，Torch-TensorRT 提供了一种导出式工作流，可以将优化后的模块序列化。这个模块可以在 PyTorch 中或与 libtorch 一起使用（即无 Python 依赖）。

第一步：优化和序列化

import torch
import torch_tensorrt

model = MyModel().eval().cuda() # 在这里定义您的模型
inputs = [torch.randn((1, 3, 224, 224)).cuda()] # 定义一个代表性输入列表

trt_gm = torch_tensorrt.compile(model, ir="dynamo", inputs=inputs)
torch_tensorrt.save(trt_gm, "trt.ep", inputs=inputs) # 对于 Python 环境来说，导出的程序仅支持 Python 运行时。C++ 部署请使用 TorchScript 文件
torch_tensorrt.save(trt_gm, "trt.ts", output_format="torchscript", inputs=inputs)

第二步：部署

在 PyTorch 中部署：

import torch
import torch_tensorrt

inputs = [torch.randn((1, 3, 224, 224)).cuda()] # 这里放置您的输入

# 可以在新的 python 会话中运行！
model = torch.export.load("trt.ep").module()
# 或者使用
# model = torch_tensorrt.load("trt.ep").module()
model(*inputs)

在 C++ 中部署：

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

auto trt_mod = torch::jit::load("trt.ts");
auto input_tensor = [...]; // 填入您的输入
auto results = trt_mod.forward({input_tensor});

支持的平台

平台	支持情况
Linux AMD64 / GPU	支持
Windows / GPU	支持（仅限 Dynamo）
Linux aarch64 / GPU	在 JetPack-4.4+ 上支持本地编译（目前请使用 v1.0.0）
Linux aarch64 / DLA	在 JetPack-4.4+ 上支持本地编译（目前请使用 v1.0.0）
Linux ppc64le / GPU	不支持

依赖

验证测试用例所使用的依赖：

Bazel 6.3.2
Libtorch 2.5.0.dev（最新 nightly 版）（使用 CUDA 12.4 编译）
CUDA 12.4
TensorRT 10.3.0.26

弃用政策

Torch-TensorRT 自版本 2.3 起，使用弃用政策通知开发人员某些 API 和工具不再推荐使用。弃用通知会在发行说明中说明，被弃用的 API 函数会在源代码中标明弃用时间。在运行时，如果使用弃用的方法和类，会发出弃用警告。Torch-TensorRT 提供 6 个月的迁移期。在此期间，API 和工具仍然可用。迁移期结束后，将以符合语义版本控制的方式移除这些 API 和工具。