預先分配輸出緩衝區

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

使用 runtime 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")