PyTorch - PyTorch - 3.5 简体中文 - UG1414

WeGO-Torch 是 WeGO 的子工程，旨在通过将 Vitis AI 工具链集成到 PyTorch 框架内以增强 Vitis AI EoU。WeGO-Torch 遵循标准 WeGO 工作流程，有助于自动进行模型分区、编译和推断，而无需人工干预。

WeGO-Torch Python/C++ API 旨在接收通过 Vitis AI PyTorch 量化器生成的量化 TorchScript 模块作为输入。然后，它会生成一个经过优化的 TorchScript 模块，可立即用于即时推断。在 PyTorch 中利用 WeGO-Torch 加速 Vitis AI 涉及以下常规步骤：

1. 将 WeGO-Torch Python 模块导入应用。
2. 使用标准 PyTorch Python API 加载 TorchScript 量化模块。
3. 使用 WeGO-Torch 模块 API 提供 TorchScript 模块和输入形状作为输入，用于编译 TorchScript 量化模块。它会返回经最优化的 Torchscript 模块作为编译结果。
4. 使用经最优化的 TorchScript 模块运行推断。

以下代码片段显示了 WeGO-Torch 的 Python API 的基本用法：

import torch
# Step 1: Import WeGO-Torch python module 
import wego_torch
 
# Step 2: Load the quantized torchscript module generated by vitis-ai PyTorch quantizer.
model_path = <quantized_torchscript_model_path>
mod = torch.jit.load(model_path)
 
# Step 3: Create an optimized TorchScript module through WeGO-Torch's API.
wego_mod = wego_torch.compile(mod,
  wego_torch.CompileOptions(
  inputs_meta = [ wego_torch.InputMeta(torch.float, [1, 3, 224, 224]) ]
  )
)
 
# Step 4: Run inference using the optmized TorchScript module.
result = wego_mod(input)

在 WeGO-Torch 中使用 C++ API 编译模型的过程与此类似。有关在 WeGO-Torch 中使用 C++ API 在 DPU 上编译和运行模型的详细信息，请参阅 核 C++ 类 和 核 C++ API。