WeGO-Torch is a sub-project of WeGO, which is designed to improve Vitis AI EoU by integrating Vitis AI toolchain into PyTorch framework. WeGO-Torch follows standard WeGO’s workflow, which means it will perform model partition, compilation and inference automatically without extra efforts from the users.
WeGO-Torch Python API will accept a quantized TorchScript module (which is generated by
Vitis AI PyTorch quantizer) as an input and returns an optimized TorchScript module,
which can be used for later inference immediately. The general steps to exploit
WeGO-Torch for Vitis-AI acceleration in PyTorch are listed below:
- Import WeGO-Torch python module into your application.
- Load the quantized TorchScript module using standard PyTorch’s Python API.
- Compile the quantized TorchScript module using wego-torch module’s API by providing the TorchScript module and input shape as inputs. It returns an optimized TorchScript module as the compiled result.
- Run inference using the optimized TorchScript module.
The following code snippets show the basic usage of Python APIs of WeGO-Torch:
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)