儲存使用 Torch-TensorRT 編譯的模型¶

可以使用 torch_tensorrt.save API 儲存使用 Torch-TensorRT 編譯的模型。

Dynamo IR¶

預設情況下，Torch-TensorRT 的 ir=dynamo 編譯輸出型別是 torch.fx.GraphModule 物件。我們可以透過指定 output_format 標誌將此物件儲存為 TorchScript (torch.jit.ScriptModule) 或 ExportedProgram (torch.export.ExportedProgram) 格式。以下是 output_format 可以接受的選項：

exported_program：這是預設選項。我們首先對 graphmodule 執行轉換，然後使用 torch.export.save 儲存模組。
torchscript：我們透過 torch.jit.trace 追蹤 graphmodule，然後透過 torch.jit.save 儲存它。

a) ExportedProgram¶

以下是一個使用示例

import torch
import torch_tensorrt

model = MyModel().eval().cuda()
inputs = [torch.randn((1, 3, 224, 224)).cuda()]
# trt_ep is a torch.fx.GraphModule object
trt_gm = torch_tensorrt.compile(model, ir="dynamo", inputs=inputs)
torch_tensorrt.save(trt_gm, "trt.ep", inputs=inputs)

# Later, you can load it and run inference
model = torch.export.load("trt.ep").module()
model(*inputs)

b) Torchscript¶

import torch
import torch_tensorrt

model = MyModel().eval().cuda()
inputs = [torch.randn((1, 3, 224, 224)).cuda()]
# trt_gm is a torch.fx.GraphModule object
trt_gm = torch_tensorrt.compile(model, ir="dynamo", inputs=inputs)
torch_tensorrt.save(trt_gm, "trt.ts", output_format="torchscript", inputs=inputs)

# Later, you can load it and run inference
model = torch.jit.load("trt.ts").cuda()
model(*inputs)

Torchscript IR¶

在 Torch-TensorRT 1.X 版本中，使用 Torchscript IR 是使用 Torch-TensorRT 編譯和執行推理的主要方式。對於 ir=ts，此行為在 2.X 版本中也保持不變。

import torch
import torch_tensorrt

model = MyModel().eval().cuda()
inputs = [torch.randn((1, 3, 224, 224)).cuda()]
trt_ts = torch_tensorrt.compile(model, ir="ts", inputs=inputs) # Output is a ScriptModule object
torch.jit.save(trt_ts, "trt_model.ts")

# Later, you can load it and run inference
model = torch.jit.load("trt_model.ts").cuda()
model(*inputs)

載入模型¶

我們可以直接使用 PyTorch 中的 torch.jit.load 和 torch.export.load API 載入 torchscript 或 exported_program 模型。此外，我們還提供了一個輕量級封裝器 torch_tensorrt.load(file_path)，它可以載入上述任一模型型別。