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.
For WeGO examples and more information on applying TensorFlow and PyTorch to deploy models, see Vitis AI GitHub repo.