使用新權重重新調整 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 允許使用者原地重新調整模組。當使用者想要更新編譯模組的權重而不建立新模組時，這非常有用。