使用新權重重新適配 Torch-TensorRT 程式

編譯是一個昂貴的操作，因為它涉及對模型進行大量的圖轉換、翻譯和最佳化。在模型權重可能偶爾更新的情況下（例如插入 LoRA 介面卡），如果每次都需要從頭開始構建編譯好的程式，則重新編譯的巨大成本可能會導致無法使用 TensorRT。Torch-TensorRT 提供了一種 PyTorch 原生機制，透過權重重適配來更新已編譯的 TensorRT 程式的權重，而無需從頭重新編譯。

在本教程中，我們將介紹

1. 將 PyTorch 模型編譯為 TensorRT 圖模組

2. 儲存和載入圖模組

3. 重新適配圖模組

本教程主要關注 AOT 工作流程，使用者在這種工作流程中很可能需要手動重新適配模組。在 JIT 工作流程中，權重變化會觸發重新編譯。由於引擎之前已經構建過，啟用引擎快取後，Torch-TensorRT 可以自動識別之前構建的引擎，觸發重新適配並代表使用者跳過重新編譯（參見：引擎快取）。

標準工作流程

匯入和模型定義

import numpy as np
import torch
import torch_tensorrt as torch_trt
import torchvision.models as models
from torch_tensorrt.dynamo import refit_module_weights

np.random.seed(0)
torch.manual_seed(0)
inputs = [torch.rand((1, 3, 224, 224)).to("cuda")]

建立可重新適配的編譯程式

初始步驟是編譯一個模組並像通常一樣儲存它。請注意，還有一個額外引數 immutable_weights 設定為 False。該引數用於指示正在構建的引擎稍後應支援權重重新適配。未設定這些引數構建的引擎將無法進行重新適配。

在此示例中，我們將編譯一個使用隨機初始化權重的 ResNet18 模型並儲存它。

model = models.resnet18(pretrained=False).eval().to("cuda")
exp_program = torch.export.export(model, tuple(inputs))
enabled_precisions = {torch.float}
debug = False
workspace_size = 20 << 30
min_block_size = 0
use_python_runtime = False
torch_executed_ops = {}
trt_gm = torch_trt.dynamo.compile(
    exp_program,
    tuple(inputs),
    use_python_runtime=use_python_runtime,
    enabled_precisions=enabled_precisions,
    debug=debug,
    min_block_size=min_block_size,
    torch_executed_ops=torch_executed_ops,
    immutable_weights=False,
    reuse_cached_engines=False,
)  # Output is a torch.fx.GraphModule

# Save the graph module as an exported program
torch_trt.save(trt_gm, "./compiled.ep", inputs=inputs)

使用預訓練權重重新適配程式

隨機權重對於推理沒有用。但現在，我們可以使用預訓練權重重新適配模型，而不是重新編譯模型。這可以透過設定另一個包含目標權重的 PyTorch 模組並將其匯出為 ExportedProgram 來完成。然後使用 refit_module_weights 函式更新已編譯模組中的權重為新權重。

# Create and compile the updated model
model2 = models.resnet18(pretrained=True).eval().to("cuda")
exp_program2 = torch.export.export(model2, tuple(inputs))


compiled_trt_ep = torch_trt.load("./compiled.ep")

# This returns a new module with updated weights
new_trt_gm = refit_module_weights(
    compiled_module=compiled_trt_ep,
    new_weight_module=exp_program2,
    arg_inputs=inputs,
)

# Check the output
expected_outputs, refitted_outputs = exp_program2.module()(*inputs), new_trt_gm(*inputs)
for expected_output, refitted_output in zip(expected_outputs, refitted_outputs):
    assert torch.allclose(
        expected_output, refitted_output, 1e-2, 1e-2
    ), "Refit Result is not correct. Refit failed"

print("Refit successfully!")

高階用法

您可以使用多種設定來控制重新適配過程

權重對映快取

權重重新適配的工作原理是將已編譯模組的權重與使用者提供的 ExportedProgram 中的新權重進行匹配。由於從 PyTorch 到 TensorRT 的 1:1 名稱匹配難以實現，因此在重新適配時匹配權重的唯一可靠方法是將新的 ExportedProgram 透過編譯過程的早期階段，以生成幾乎相同的權重名稱。這可能很昂貴，並非總是必要。

為了避免這種情況，在初始編譯時，Torch-TensorRT 將嘗試快取從 PyTorch 權重到 TensorRT 權重的直接對映。此快取作為元資料儲存在已編譯模組中，可用於加快重新適配速度。如果快取不存在，重新適配系統將回退到在重新適配時重建對映。此快取的使用由 use_weight_map_cache 引數控制。

由於快取使用基於啟發式的方法來匹配 PyTorch 和 TensorRT 權重，您可能需要驗證重新適配。這可以透過將 verify_output 設定為 `True` 並提供示例 arg_inputs 和 kwarg_inputs 來完成。完成此操作後，重新適配系統將在相同的輸入上執行重新適配後的模組和使用者提供的模組，並比較輸出。

原地重新適配

in_place 引數允許使用者在原地重新適配模組。當用戶想要更新已編譯模組的權重而無需建立新模組時，此功能非常有用。