WiGNet: Windowed Vision Graph Neural Network

TL;DR

A novel Windowed vision Graph neural Network (WiGNet) model for efficient image processing that uses graph convolutions instead of the typical 2D convolution or self-attention mechanism and offers a promising solution toward the deployment of vision GNNs in real-world applications.

Abstract

In recent years, Graph Neural Networks (GNNs) have demonstrated strong adaptability to various real-world challenges, with architectures such as Vision GNN (ViG) achieving state-of-the-art performance in several computer vision tasks. However, their practical applicability is hindered by the computational complexity of constructing the graph, which scales quadratically with the image size. In this paper, we introduce a novel Windowed vision Graph neural Network (WiGNet) model for efficient image processing. WiGNet explores a different strategy from previous works by partitioning the image into windows and constructing a graph within each window. Therefore, our model uses graph convolutions instead of the typical 2D convolution or self-attention mechanism. WiGNet effectively manages computational and memory complexity for large image sizes. We evaluate our method in the ImageNet-1k benchmark dataset and test the adaptability of WiGNet using the CelebA-HQ dataset as a downstream task with higher-resolution images. In both of these scenarios, our method achieves competitive results compared to previous vision GNNs while keeping memory and computational complexity at bay. WiGNet offers a promising solution toward the deployment of vision GNNs in real-world applications. We publicly released the code at https://github.com/EIDOSLAB/WiGNet.

Figures (6)

• Figure 1: Implementation of Vision GNN (ViG) vgnn and the proposed WiGNet. Our model first divides images into local windows where graphs are built. This fundamental change dramatically increases computational and memory efficiency in vision tasks.
• Figure 2: (a): WiGNet architecture exemplified for the Tiny version (see Table \ref{['tab:wignn']} for details). In this example, a linear classifier generates class scores. (b): A graphical illustration of the Window-based Grapher module.
• Figure 3: Overview of the cycling operation used to obtain shifted windows. The top-left part of the feature maps is copied on the bottom-right part, then the masking mechanism is used to avoid connection between non-adjacent nodes in the original feature maps.
• Figure 4: Illustrative example of the dynamic graph convolution of WiGNet.
• Figure 5: Computational complexity and GPU memory footprint of several vision GNN architectures and WiGNet in terms of MACs and MB on NVIDIA GeForce RTX 3090 GPU.