Using WeGO - 3.5 English

Vitis AI User Guide (UG1414)

Document ID
Release Date
3.5 English

WeGO (Whole Graph Optimizer) is a key feature of Vitis AI, providing a seamless solution for deploying TensorFlow and PyTorch models on the Versal Data Center DPU. By integrating the Vitis AI development kit with TensorFlow and PyTorch frameworks, WeGO enables efficient deployment of models on the DPU. WeGO supports TensorFlow 2.x and PyTorch 1.x, starting from Vitis AI 2.5. It is considered an in-framework inference solution within Vitis AI, setting it apart from non-framework approaches that use VART APIs or the AI Library.

WeGO automatically performs subgraph partitioning and applies optimizations and acceleration for the data center DPU-compatible subgraphs. The parts of the graph that are not supported by the DPU, known as CPU subgraphs, are dispatched to the TensorFlow or PyTorch framework for CPU native execution. WeGO handles the entire process, including whole graph optimization, compilation, and runtime subgraphs' dispatch and execution. This process remains entirely transparent to end-users, ensuring ease of use.

Using WeGO offers model designers a seamless transition from training to inference. With WeGO, deploying quantized models over a TensorFlow or PyTorch framework becomes straightforward through its Python programming interface. This interface enables maximum reuse of Python code, including pre-processing and post-processing, developed during the model training phase with TensorFlow or PyTorch. As a result, WeGO significantly enhances productivity when deploying and evaluating models over data center DPUs.
Note: Currently, WeGO only supports data center DPU target DPUCV2DX8G on the V70 platform.

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