預分配輸出緩衝區

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