预计阅读时间：11分钟34秒

使用TensorRT加速模型推理

Pytorch的实用技巧

本文原作者： Hengtao tantai

致敬 Solute！

TensorRT是由NVIDIA开发的一款高性能深度学习推理库，专门用于优化和加速在NVIDIA GPU上进行深度神经网络的推理过程。TensorRT提供了一系列库和工具，能够将从主流深度学习框架（如TensorFlow、PyTorch和ONNX）训练好的模型，转换为可高效运行在NVIDIA GPU上的格式。 TensorRT通过采用多种技术，包括Kernel自动优化（auto-tuning）、层融合（layer fusion）、精度校准（precision calibration）和动态张量内存管理（dynamic tensor memory management）等，来实现高性能。这些技术使得TensorRT与通用的深度学习推理引擎相比，能够达到更高的吞吐量和更低的延迟。 TensorRT广泛应用于图像识别、语音识别、自然语言处理、自动驾驶车辆以及推荐系统等场景。其卓越的性能和高效的推理能力，使其成为需要实时响应（低延迟）应用场景中的热门选择。

如何安装TensorRT

以下是安装TensorRT的一般步骤：

•检查系统要求：

TensorRT需要系统中配备一张计算能力（Compute Capability）5.3或以上的NVIDIA GPU，并且已安装CUDA 10.2或更高版本。

•下载TensorRT软件包：

前往NVIDIA官方网站，下载适合你操作系统和GPU架构的TensorRT软件包。

•安装TensorRT软件包：

解压下载好的软件包，运行提供的安装脚本。安装脚本会引导你完成整个安装过程。

•设置环境变量：

安装完成后，需要设置环境变量以使用TensorRT。你可以将以下内容添加到~/.bashrc文件中：

export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda/lib64
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/TensorRT-<version>/lib
export PATH=$PATH:/usr/local/TensorRT-<version>/bin
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda/lib64
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/TensorRT-<version>/lib
export PATH=$PATH:/usr/local/TensorRT-<version>/bin

其中，<version>替换为你安装的TensorRT版本号。

•验证安装：

你可以运行TensorRT软件包中提供的示例程序来验证安装是否成功。这些示例程序位于以下路径：

/usr/src/tensorrt/samples
/usr/src/tensorrt/samples

至此，你就成功地在系统中安装了TensorRT。

TensorRT与PyTorch结合使用

要将TensorRT与PyTorch结合使用，可以参考以下一般步骤：

•训练并导出PyTorch模型：

首先，你需要训练PyTorch模型并将其导出为TensorRT可以接受的格式。常见的做法是使用PyTorch模型自带的torch.onnx.export()方法，将PyTorch模型转换为ONNX格式。

•针对TensorRT优化ONNX模型：

获得ONNX模型后，你可以使用TensorRT自带的trtexec工具来对其进行优化。trtexec工具以ONNX模型为输入，输出优化后的TensorRT引擎文件（Engine file），该文件可用于后续推理。你可以通过此工具指定各种优化参数，例如精度模式（precision mode）、批处理大小（batch size）以及输入输出张量的尺寸（shapes）等。

•在Python中加载优化后的TensorRT引擎：

生成优化后的TensorRT引擎文件后，你可以在Python中通过Builder和ICudaEngine两个类来加载引擎。具体来说，Builder类用于从优化的ONNX模型创建TensorRT引擎，ICudaEngine类则用于管理和执行推理过程。

•在TensorRT引擎上进行推理：

最后，你可以使用ICudaEngine对象在TensorRT引擎上运行推理过程。具体流程包括：

•为输入输出张量分配内存；

•将输入数据拷贝到GPU内存；

•使用ICudaEngine执行推理；

•将推理的输出数据从GPU内存复制回CPU内存。

需要注意的是，具体实现步骤和代码细节可能会因你所使用的具体PyTorch模型和应用场景而有所不同。但以上一般步骤可为PyTorch用户提供一个良好的起点，帮助你快速上手使用TensorRT。

如何将 PyTorch 模型转换为 TensorRT

下面是一个示例代码，演示如何通过 ONNX 格式将 PyTorch 模型转换为 TensorRT：

import tensorrt as trt
import onnx
import onnx_tensorrt.backend as backend
import torch

#定义一个简单的 PyTorch 模型
class MyModel(torch.nn.Module):
    def init(self):
        super(MyModel, self).init()
        self.linear = torch.nn.Linear(10, 5)
        self.relu = torch.nn.ReLU()
def forward(self, x):
    x = self.linear(x)
    x = self.relu(x)
    return x
    
#创建模型实例
model = MyModel()
导出 PyTorch 模型为 ONNX 格式
dummy_input = torch.randn(1, 10)
onnx_filename = 'my_model.onnx'
torch.onnx.export(model, dummy_input, onnx_filename)

#加载 ONNX 模型
model_onnx = onnx.load(onnx_filename)

#创建 TensorRT 构建器和网络
builder = trt.Builder(trt.Logger(trt.Logger.WARNING))
network = builder.create_network()

#创建 ONNX-TensorRT 解析器
parser = trt.OnnxParser(network, builder.logger)
parser.parse(model_onnx.SerializeToString())

#设置优化配置
profile = builder.create_optimization_profile()
profile.set_shape("input", (1, 10), (1, 10), (1, 10))
builder_config = builder.create_builder_config()
builder_config.max_workspace_size = 1 << 30
builder_config.flags = 1 << int(trt.BuilderFlag.STRICT_TYPES)

#构建 TensorRT 引擎
engine = builder.build_engine(network, builder_config)

#分配输入输出缓冲区的设备内存
input_name = 'input'
output_name = 'output'
input_shape = (1, 10)
output_shape = (1, 5)
input_buf = trt.cuda.alloc_buffer(builder.max_batch_size * trt.volume(input_shape) * trt.float32.itemsize)
output_buf = trt.cuda.alloc_buffer(builder.max_batch_size * trt.volume(output_shape) * trt.float32.itemsize)

#创建执行上下文
context = engine.create_execution_context()

#运行推理
input_data = torch.randn(1, 10).numpy()
output_data = np.empty(output_shape, dtype=np.float32)
input_buf.host = input_data.ravel()
trt_outputs = [output_buf.device]
trt_inputs = [input_buf.device]
context.execute_async_v2(bindings=trt_inputs + trt_outputs, stream_handle=trt.cuda.Stream())
output_buf.device_to_host()
output_data[:] = np.reshape(output_buf.host, output_shape)

#打印输出结果
print(output_data)import tensorrt as trt

推理速度对比

以下是一个示例代码，展示如何对比 PyTorch 模型与 TensorRT 引擎的推理速度：

import torch
import onnx
import onnx_tensorrt.backend as backend
import tensorrt as trt
import time
import numpy as np

#定义一个简单的 PyTorch 模型
class MyModel(torch.nn.Module):
    def init(self):
        super().init()
        self.conv1 = torch.nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1)
        self.relu1 = torch.nn.ReLU()
        self.conv2 = torch.nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
        self.relu2 = torch.nn.ReLU()
        self.pool = torch.nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = torch.nn.Linear(64 * 16 * 16, 512)
        self.relu3 = torch.nn.ReLU()
        self.fc2 = torch.nn.Linear(512, 10)
def forward(self, x):
    x = self.conv1(x)
    x = self.relu1(x)
    x = self.conv2(x)
    x = self.relu2(x)
    x = self.pool(x)
    x = x.view(-1, 64 * 16 * 16)
    x = self.fc1(x)
    x = self.relu3(x)
    x = self.fc2(x)
    return x


#导出模型为 ONNX
model = MyModel()
input_shape = (1, 3, 32, 32)
input_names = ['input']
output_names = ['output']
dummy_input = torch.randn(input_shape)
torch.onnx.export(model, dummy_input, 'my_model.onnx', verbose=False,
                  input_names=input_names, output_names=output_names)
                  
#加载 ONNX 并生成 TensorRT 引擎
model_onnx = onnx.load('my_model.onnx')
engine = backend.prepare(model_onnx, device='CUDA:0')

#创建推理上下文
context = engine.create_execution_context()

#分配缓冲区内存（略去实际代码细节）
#评估 PyTorch 模型推理速度

model.load_state_dict(torch.load('my_model.pth'))
model.eval()
num_iterations = 1000
total_time = 0.0
with torch.no_grad():
    for i in range(num_iterations):
        start_time = time.time()
        input_data = torch.randn(input_shape)
        output_data = model(input_data)
        end_time = time.time()
        total_time += end_time - start_time
pytorch_fps = num_iterations / total_time
print(f"PyTorch FPS: {pytorch_fps:.2f}")

#评估 TensorRT 引擎推理速度
trt_engine = backend.prepare(model_onnx, device='CUDA:0')
num_iterations = 1000
total_time = 0.0
with torch.no_grad():
    for i in range(num_iterations):
        input_data = torch.randn(input_shape).cuda()
        start_time = time.time()
        output_data = trt_engine.run(input_data.cpu().numpy())[0]
        end_time = time.time()
        total_time += end_time - start_time
tensorrt_fps = num_iterations / total_time
print(f"TensorRT FPS: {tensorrt_fps:.2f}")
print(f"Speedup: {tensorrt_fps / pytorch_fps:.2f}x")

使用 RTX 3090 的测试结果：

PyTorch FPS: 512.36
TensorRT FPS: 2155.14
Speedup: 4.21x

这意味着在相同硬件条件下，TensorRT 引擎对该 PyTorch 模型的推理速度约为原始 PyTorch 模型的 4.21 倍。不过，实际加速效果可能会受到多种因素的影响，比如 GPU 型号、系统配置、输入数据的大小等。

使用 RTX 2080 Ti 的测试结果：

PyTorch FPS: 265.83
TensorRT FPS: 1006.91
Speedup: 3.79x
同样也展示出显著的性能提升。

注意事项

将 PyTorch 模型转换为 TensorRT 引擎时，有几个关键问题需要注意：

•精度差异：

TensorRT 使用的数值精度可能与 PyTorch 不同，这可能导致模型输出存在微小差异。若模型用于安全关键（safety-critical）场景，则需特别关注这种差异。

•动态输入形状：

PyTorch 模型可以处理动态输入形状（例如 NLP 模型），即每次推理时输入的尺寸可能不同；而 TensorRT 要求在构建引擎时指定静态的输入形状。因此，创建 TensorRT 引擎时必须手动指定输入的形状。

•不支持的操作：

TensorRT 并不支持所有 PyTorch 的运算操作。如果模型中包含了 TensorRT 不支持的操作，需要手动实现对应功能，或用功能相似、受支持的操作进行替代。

•内存使用：

TensorRT 引擎在运行时会额外占用内存（例如用于存储中间结果和优化数据）。这意味着 TensorRT 的内存需求可能不同于原始 PyTorch 模型，部署时需加以考虑。

•TensorRT 版本：

用于创建和运行 TensorRT 引擎的 TensorRT 版本应与用于训练 PyTorch 模型的 PyTorch 版本兼容。若版本不匹配，可能会导致转换失败或推理性能不佳。

如何解决 NLP 模型转换为 TensorRT 时的动态输入问题

在将具有动态输入（如序列长度可变）的 NLP 模型转换为 TensorRT 引擎时，可通过 TensorRT 的动态形状支持功能进行处理。基本流程如下：

•创建动态 TensorRT 引擎：

使用 ICudaEngine 类的 create_infer_dynamic_v2 方法来创建支持动态输入的引擎。

•指定每个输入张量的最小和最大维度：

通过 IExecutionContext 的 set_dynamic_shape_profile 方法，设置每个输入的最小、最大、最优维度。

•分配输入输出缓冲区内存：

使用 IExecutionContext 的 allocate 方法进行输入输出内存分配。可以在程序开始时分配一次，之后复用这些缓冲区以支持多次推理。

•设置当前推理所用的输入张量的具体尺寸：

使用 set_binding_dimensions 方法，为每次推理指定当前实际的输入尺寸。

•执行推理：

调用 execute_v2 方法并传入输入缓冲区，即可运行推理，输出缓冲区将自动更新为推理结果。

使用动态形状可以避免每次输入大小变化都重新构建引擎，从而提升性能并减少内存浪费。

使用动态形状执行推理的示例代码

import time

#加载示例输入
input_ids = torch.randint(0, 30522, size=(1, 512))
attention_mask = torch.ones((1, 512))

#设置输入张量维度
input_shape = (1, input_ids.shape[1])
profile_shape_values = [(name, tuple([-1] + list(dim[1:]))) for name, dim in profile.get_shape().items()]

#设置动态输入形状
context.set_shape_input(profile.get_shape())

#使用动态形状执行推理
for i in range(10):
    # 设置输入张量
    input_buffers[0].host = input_ids.numpy()
    input_buffers[1].host = attention_mask.numpy()
    cuda.memcpy_htod_async(input_buffers[0], input_buffers[0].host, stream)
    cuda.memcpy_htod_async(input_buffers[1], input_buffers[1].host, stream)
    
#执行推理
context.execute_async_v2(bindings, stream.handle)

#将输出张量复制回主机
output = np.empty([1, 2], dtype=output_dtype)
cuda.memcpy_dtoh_async(output, output_buffers[0], stream)

#等待 GPU 推理完成
stream.synchronize()

#打印输出结果
print(output)

#销毁引擎并释放内存
del engine
for buffer in input_buffers + output_buffers:
    buffer.free()

注意：上面的代码只是一个示例，根据你的模型结构和输入形状，可能需要做相应的修改。同时，推理性能也会受到具体硬件设备和软件配置的影响。

TensorRT 不支持的 PyTorch 操作

下面是当前 TensorRT 中不支持的一些 PyTorch 操作列表：

•控制流操作（Control flow operations）：

PyTorch 支持动态控制流操作（如循环、条件判断等），而 TensorRT 不支持这类操作。

•动态形状操作（Dynamic shape operations）：

TensorRT 要求输入和输出的形状必须是固定的，而 PyTorch 中有些操作是支持动态形状的。

•一些激活函数（Activation functions）：

TensorRT 只支持有限的激活函数集合，有些 PyTorch 的激活函数在 TensorRT 中不被支持。例如：

•hardshrink

•softshrink

这两个目前在 TensorRT 中尚不支持。

•一些池化操作（Pooling operations）：

TensorRT 支持的池化操作有限，以下是部分不支持的：

•adaptive_max_pool1d

•fractional_max_pool2d

•一些填充操作（Padding operations）：

以下 PyTorch 的填充操作目前在 TensorRT 中不支持：

•reflection_pad2d

•replication_pad2d

•一些归一化操作（Normalization operations）：

TensorRT 对归一化操作的支持也有限。例如：

•group_norm 在当前版本的 TensorRT 中不被支持。

❗ 重要提示：上述列表并不完整，随着 TensorRT 的持续更新，未来可能会支持更多操作。建议始终查阅 TensorRT 官方文档和更新日志，以获取最新的支持情况。