预分配的输出缓冲区

TensorRT 运行时模块充当一个包装器，围绕着一个已经被编译并优化为 TensorRT 引擎的 PyTorch 模型（或子图）。

当编译的模块被执行时，输入和输出张量被设置为TensorRT上下文进行处理。 如果在TensorRT上下文执行后将输出缓冲区分配并用于下一次推理，GPU任务和内存分配任务可以并发操作。这种重叠允许更有效地使用GPU资源，可能提高推理的性能。

这种优化在以下情况下特别有效

1. Small inference time

   • 输出缓冲区的分配通常需要最少的CPU周期，因为缓存机制有效地处理了内存重用。与整体推理时间相比，这种分配所需的时间相对恒定，从而带来了显著的性能提升，特别是在涉及小型推理工作负载的场景中。这是因为当计算工作负载不足以掩盖这些节省时，减少的分配时间有助于加快执行速度。

2. Multiple graph breaks

   • 如果模块包含TensorRT不支持的操作，不支持的部分将由PyTorch处理，这种回退会导致图形中断。跨多个子图的优化缓冲区分配的累积效应可以提高整体推理性能。

   • 虽然优化输出缓冲区可以减轻部分开销，但应优先考虑减少或消除图形中断，因为它能够实现更全面的优化

3. Static input or infrequent input shape change

   • 如果形状发生变化，预分配的缓冲区不能用于下一次推理，并且在执行TensorRT上下文之前会有新的分配。此功能不适用于输入形状频繁变化的用例

导入和模型定义

import timeit

import numpy as np
import torch
import torch_tensorrt
from transformers import BertModel

定义函数以衡量推理性能

def test_module_perf(model, *input):
    timings = []

    # Warm-up phase to ensure consistent and accurate performance measurements.
    with torch.no_grad():
        for _ in range(3):
            model(*input)
    torch.cuda.synchronize()

    # Timing phase to measure inference performance
    with torch.no_grad():
        for i in range(10):
            start_time = timeit.default_timer()
            model(*input)
            torch.cuda.synchronize()
            end_time = timeit.default_timer()
            timings.append(end_time - start_time)
    times = np.array(timings)
    time_med = np.median(times)

    # Return the median time as a representative performance metric
    return time_med

加载模型并编译

# Load bert model
model = (
    BertModel.from_pretrained("bert-base-uncased", torchscript=True)
    .eval()
    .half()
    .to("cuda")
)
# Define sample inputs
inputs = [
    torch.randint(0, 5, (1, 128), dtype=torch.int32).to("cuda"),
    torch.randint(0, 5, (1, 128), dtype=torch.int32).to("cuda"),
]
# Next, we compile the model using torch_tensorrt.compile
optimized_model = torch_tensorrt.compile(
    model,
    ir="dynamo",
    enabled_precisions={torch.half},
    inputs=inputs,
)

使用运行时API启用/禁用预分配输出缓冲区功能

# Enable pre-allocated output buffer using a context manager
with torch_tensorrt.runtime.enable_pre_allocated_outputs(optimized_model):
    out_trt = optimized_model(*inputs)
    # Subsequent inferences can use the pre-allocated output buffer (no shape change)
    out_trt = optimized_model(*inputs)

# Alternatively, we can enable the feature using a context object
pre_allocated_output_ctx = torch_tensorrt.runtime.enable_pre_allocated_outputs(
    optimized_model
)
pre_allocated_output_ctx.set_pre_allocated_output(True)
time_opt = test_module_perf(optimized_model, *inputs)

# Disable the pre-allocated output buffer feature and perform inference normally
pre_allocated_output_ctx.set_pre_allocated_output(False)
out_trt = optimized_model(*inputs)
time_normal = test_module_perf(optimized_model, *inputs)

time_opt_ms = time_opt * 1000
time_normal_ms = time_normal * 1000

print(f"normal trt model time: {time_normal_ms:.3f} ms")
print(f"pre-allocated output buffer model time: {time_opt_ms:.3f} ms")