Using WeGO - 3.0 English

Vitis AI User Guide (UG1414)

Document ID
Release Date
3.0 English

WeGO (Whole Graph Optimizer) is a Vitis AI early access feature that offers a smooth solution to deploy TensorFlow and PyTorch models on cloud DPU by integrating Vitis AI development kit with TensorFlow and PyTorch frameworks. In addition to TensorFlow 1.15, WeGO begins to support TensorFlow 2.8 and PyTorch 1.10 since VAI 2.5. WeGO is regarded as Vitis AI in-framework inference solution, compared with non-framework approach using VART APIs or AI Library.

WeGO automatically performs subgraph partitioning for the Vitis AI quantized models and applies optimizations and acceleration for the cloud DPU compatible subgraphs. The remaining DPU unsupported parts of the graph, called CPU subgraphs, are dispatched to TensorFlow or PyTorch framework for CPU native execution. WeGO takes care of the whole graph optimization, compilation, and run-time subgraphs’ dispatch and execution. This process is entirely transparent to the end-users, making it easy to use.

Using WeGO is a straightforward transition from training to inference for model designers. WeGO provides a Python programming interface to deploy the quantized models over the TensorFlow or PyTorch framework. This makes it possible to maximally reuse the Python code (including pre-processing and post-processing) developed during models training phase with TensorFlow or PyTorch. It substantially improves the productivity for models’ deployment and evaluation over cloud DPUs.
Note: Currently, WeGO only supports cloud DPU target DPUCVDX8H on VCK5000 Prod platform.

For WeGO examples and more information on how to apply TensorFlow and PyTorch to deploy models, see Vitis AI GitHub repo.