预分配输出缓冲区

TensorRT 运行时模块是 PyTorch 模型（或子图）的包装器，该模型已编译并优化为 TensorRT 引擎。

执行编译后的模块时，输入和输出张量会被设置到 TensorRT 上下文中进行处理。如果将输出缓冲区的分配移至 TensorRT 上下文执行之后，并将其用于下一次推理，则 GPU 任务和内存分配任务可以并行运行。这种重叠可以更高效地利用 GPU 资源，从而潜在地提高推理性能。

此优化在以下情况下尤其有效

1. 推理时间较短

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

2. 多个图分割

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

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

3. 静态输入或输入形状不常改变

   • 如果形状发生改变，预分配的缓冲区将无法用于下一次推理，并且在执行 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")