预分配输出缓冲区

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")