ViG-UNet: Vision Graph Neural Networks for Medical Image Segmentation
Deep neural networks have been widely used in medical image analysis and
medical image segmentation is one of the most important tasks. U-shaped neural
networks with encoder-decoder are prevailing and have succeeded greatly in
various segmentation tasks. While CNNs treat an image as a grid of pixels in
Euclidean space and Transformers recognize an image as a sequence of patches,
graph-based representation is more generalized and can construct connections
for each part of an image. In this paper, we propose a novel ViG-UNet, a graph
neural network-based U-shaped architecture with the encoder, the decoder, the
bottleneck, and skip connections. The downsampling and upsampling modules are
also carefully designed. The experimental results on ISIC 2016, ISIC 2017 and
Kvasir-SEG datasets demonstrate that our proposed architecture outperforms most
existing classic and state-of-the-art U-shaped networks.
Juntao Jiang, Xiyu Chen, Guanzhong Tian, Yong Liu
# Vig-UNet: Vision Graph Neural Networks for Medical Image Segmentation

###### Abstract

Deep neural networks have been widely used in medical image analysis and medical image segmentation is one of the most important tasks. U-shaped neural networks with encoder-decoder are prevailing and have succeeded greatly in various segmentation tasks. While CNNs treat an image as a grid of pixels in Euclidean space and Transformers recognize an image as a sequence of patches, graph-based representation is more generalized and can construct connections for each part of an image. In this paper, we propose a novel ViG-UNet, a graph neural network-based U-shaped architecture with the encoder, the decoder, the bottleneck, and skip connections. The downsampling and upsampling modules are also carefully designed. The experimental results on ISIC 2016, ISIC 2017 and Kvasir-SEG datasets demonstrate that our proposed architecture outperforms most existing classic and state-of-the-art U-shaped networks.

Juntao Jiang\({}^{1}\), Xiyu Chen\({}^{2}\), Guanzhong Tian\({}^{3}\)1 and Yong Liu 1\({}^{*}\)1. College of Control Science and Engineering, Zhejiang University, Hangzhou, China

2. Polytechnic Institute, Zhejiang University, Hangzhou, China

3. Ningbo Innovation Center, Zhejiang University, Ningbo, China Medical image segmentation, ViG-UNet, Graph neural networks, Encoder-decoder

Footnote 1: Corresponding authors

## 1 Introduction

Recent years have witnessed the rise of deep learning and its broader applications in computer vision tasks. As one of the most heated topics of computer vision applied in medical scenarios, image segmentation, identifying the pixels of organs or lesions from the background, plays a crucial role in computer-aided diagnosis and treatment, improving efficiency and accuracy.

Currently, medical image segmentation methods based on deep learning mainly use fully convolutional neural networks (FCN) with U-shaped encoder-decoder architecture such as U-Net [1] and its variants. Composed of a symmetric encoder-decoder with skip connections, U-Net uses convolutional layers and downsampling modules for feature extraction, while convolutional layers and upsampling modules for pixel-level semantic classification. The skip connection operation can maintain spatial information from a high-resolution feature, which may be lost in downsampling. Following this work and based on a fully convolutional structure, a lot of U-Net's variants like Attention-UNet [2], UNet++ [3] and so on, have been proposed and achieved great success. Recently, as Transformer-based methods like ViT [4] achieved good results in image recognition tasks, thanks to their capability of enhancing global understanding of images, extracting information from the inputs and their interrelations, the Transformer-based medical image segmentation models such as Trans-UNet [5] and Swin-UNet [6] also have been proposed and showed competitive performance.

While CNNs treat an image as a grid of pixels in Euclidean space and Transformer recognizes an image as a sequence of patches, graph-based representation can be more generalized and reflect the relationship of each part in an image. Since the graph neural network (GNN) [7] was first proposed, the techniques for processing graphs have been researched a lot. A series of spatial-based GCNs [8, 9] and spectral-based GCNs [10, 11, 12, 13] are widely proposed and applied. In recent work, Han et al. [14] proposed a Vision GNN (ViG), which splits the image into many blocks regarded as nodes and constructs a graph representation by connecting the nearest neighbors, then uses GNNs to process it. It contains Grapher modules with graph convolution to aggregate and update graph information and Feed-forward Networks (FFNs) modules with fully connected networks for node feature transformation, which performed well in image recognition tasks. ViG-S has achieved 0.804 Top-1 accuracy and ViG-B has achieved 0.823 on ImageNet [15].

Motivated by the success of ViG model, we propose a ViG-UNet to utilize the powerful functions of ViG for 2D medical image segmentation in this work. The graph-based representation can also be effective in segmentation tasks. ViG-UNet is a GNN-based U-shaped architecture consisting of the encoder, the bottleneck and the decoder, with skip connections. We do comparison experiments on ISIC 2016 [16], ISIC 2017 [17] and Kvair-SEG [18] datasets.The results show that the proposed model outperformed most existing classic and state-of-the-art methods. The code will be released at _[https://github.com/juntaoJianggavin/ViG-UNet_](https://github.com/juntaoJianggavin/ViG-UNet_).
## 2 Methods

### Architecture Overview

ViG-UNet is a U-shape model with symmetrical architectures, whose architecture can be seen in Figure 1. It consists of structures of the encoder, the bottleneck, the decoder and skip-connections. The basic modules of ViG-UNet are the stem block, Grapher Modules, Feed-forward Networks (FFNs), downsampling and upsampling modules. The detailed settings of ViG-UNet can be seen in Table 1, where \(D\) means feature dimension, \(E\) means the numbers of convolutional layers in FFNs, \(K\) means the number of neighbors in GCNs, \(H\times W\) means the output image size.

### Stem Block, Upsampling, Downsampling and SkipConnections

In the stem block, two convolutional layers are used with stride 1 and stride 2, respectively. The output features have height and width equal to \(\frac{H}{2}\) and \(\frac{W}{2}\), where \(H\), \(W\) are the original height and width of the input image. And the position embedding is added. We used a convolutional layer with stride 2 for downsampling operation and a bilinear for upsampling operation with the scale factor 2 following with a convolutional layer. The output of each FFN in the encoder is added to the output of the FFN in the decoder.

### Grapher Module

Vision GNN first builds an image's graph structure by dividing it into \(N\) patches, converting them into feature vectors, and then recognizing them as a set of nodes \(\mathcal{V}=\{v_{1},v_{2},\cdots,v_{N}\}\). A \(K\) nearest neighbors method is used to find \(K\) nearest neighbors \(\mathcal{N}\left(v_{i}\right)\) for each node \(v_{i}\). An edge \(e_{ji}\) is added from \(v_{j}\) to \(v_{i}\) for all \(v_{j}\in\mathcal{N}\left(v_{i}\right)\). In this way, a graph representation of an image \(\mathcal{G}=\left(\mathcal{V},\mathcal{E}\right)\) is obtained, where \(e_{ji}\in\mathcal{E}\). For the constructed graph representation \(\mathcal{G}=G(X)\) and the input feature \(x_{i}\), the aggregation operation calculates the representation of a node by aggregating features of neighboring nodes. Then the update operation merge the aggregated feature. The updated feature \(\mathbf{x}_{i}^{\prime}\) can be represented as:

\[\mathbf{x}_{i}^{\prime}=h\left(\mathbf{x}_{i},g\left(\mathbf{x}_{i},\mathcal{ N}\left(\mathbf{x}_{i}\right);W_{aggregate}\right);W_{\text{update}}\right), \tag{1}\]

where \(W_{aggregate}\) and \(W_{update}\) are the learnable weights of the aggregation and update operations.

\[g(\cdot)=\mathbf{x}_{i}^{\prime\prime}=\left[\mathbf{x}_{i},\max\left(\left\{ \mathbf{x}_{j}-\mathbf{x}_{i}\mid j\in\mathcal{N}\left(\mathbf{x}_{i}\right) \right\}\right]\right., \tag{2}\]

and

\[h(\cdot)=\mathbf{x}_{i}^{\prime}=\mathbf{x}_{i}^{\prime\prime}W_{\text{ update}}+b_{h}, \tag{3}\]

\(b_{h}\) is the bias. And following the design of original ViG networks, the \(g(\cdot)\) operation uses the max-relative graph convolution [19].

\begin{table}
\begin{tabular}{c|c|c} \hline Module & Output size & Architecture \\ \hline Stem & \(\frac{H}{2}\times\frac{W}{2}\) & Conv\(\times\)2 \\ Grapher + FFN & \(\frac{H}{2}\times\frac{W}{2}\) & \(\left[\begin{array}{c}D=32\\ E=2\\ K=9\end{array}\right]\) \\ Downsampling & \(\frac{H}{4}\times\frac{W}{4}\) & Conv \\ Grapher + FFN & \(\frac{H}{4}\times\frac{W}{4}\) & \(\left[\begin{array}{c}D=64\\ E=2\\ K=9\end{array}\right]\) \\ Downsampling & \(\frac{H}{8}\times\frac{W}{8}\) & Conv \\ Grapher + FFN & \(\frac{H}{8}\times\frac{W}{8}\) & \(\left[\begin{array}{c}D=128\\ E=2\\ K=9\end{array}\right]\) \\ Downsampling & \(\frac{H}{16}\times\frac{W}{16}\) & Conv \\ Grapher + FFN & \(\frac{H}{16}\times\frac{W}{16}\) & \(\left[\begin{array}{c}D=256\\ E=2\\ K=9\end{array}\right]\) \\ Downsampling & \(\frac{H}{32}\times\frac{W}{32}\) & Conv \\ Grapher \(\times\)2 & \(\frac{H}{32}\times\frac{W}{32}\) & \(\left[\begin{array}{c}D=512\\ K=9\end{array}\right]\times 2\) \\ Upsampling & \(\frac{H}{16}\times\frac{W}{16}\) & bilinear + Conv \\ Grapher + FFN & \(\frac{H}{16}\times\frac{W}{16}\) & \(\left[\begin{array}{c}D=256\\ E=2\\ K=9\end{array}\right]\) \\ Upsampling & \(\frac{H}{8}\times\frac{W}{8}\) & bilinear + Conv \\ Grapher + FFN & \(\frac{H}{8}\times\frac{W}{8}\) & \(\left[\begin{array}{c}D=128\\ E=2\\ K=9\end{array}\right]\) \\ Upsampling & \(\frac{H}{4}\times\frac{W}{4}\) & bilinear + Conv \\ Grapher + FFN & \(\frac{H}{4}\times\frac{W}{4}\) & \(\left[\begin{array}{c}D=64\\ E=2\\ K=9\end{array}\right]\) \\ Upsampling & \(\frac{H}{2}\times\frac{W}{2}\) & bilinear + Conv \\ Grapher + FFN & \(\frac{H}{2}\times\frac{W}{2}\) & \(\left[\begin{array}{c}D=32\\ E=2\\ K=9\end{array}\right]\) \\ Final Layer & \(H\times W\) & bilinear + Conv \\ \hline \end{tabular}
\end{table}
Table 1: Detailed settings of ViG-UNet
The aggregated feature \(\mathbf{x}_{i}^{\prime\prime}\) is split into \(h\) heads, then each head is updated with different weights. Then the updated feature \(\mathbf{x}_{i}^{\prime}\) can be obtained by concatenating all the heads:

\[\begin{split}\mathbf{x}_{i}^{\prime}=[\mathbf{x}_{\text{ineal}}^{ \prime\prime}\,\,W_{\text{update}}^{1}\,+b_{h1},\mathbf{x}_{\text{ineal}\,2}^{ \prime\prime}\,W_{\text{update}}^{2}\,+b_{h2},\\ \mathbf{x}_{\text{ineal}\,3}^{\prime\prime}\,W_{\text{update}}^{3} \,+b_{h3},\cdots\mathbf{x}_{\text{ineal}h}^{\prime\prime}\,W_{\text{update}}^{ h}\,+b_{hh}]\end{split} \tag{4}\]

where \(\mathbf{x}_{\text{ineal}\,1}^{\prime\prime}\,\mathbf{x}_{\text{ineal}\,2}^{ \prime\prime}\,\cdots,\mathbf{x}_{\text{ineal}h}^{\prime\prime}\) represent the split heads from \(\mathbf{x}_{i}^{\prime}\), \(\,W_{\text{update}}^{1}\,W_{\text{update}}^{2}\,\cdots,W_{\text{update}}^{h}\,\) represent different weights and \(b_{h1},b_{h2},\cdots,b_{hh}\) represent different biases.

For the input feature \(X\), the output feature \(Y\) after a Grapher module can be represented as:

\[X_{1}=XW_{in}+b_{in}, \tag{5}\]

\[Y=\mathrm{Droppath}(\mathrm{GELU}(\mathrm{GraphConv}(X_{1})W_{\text{out}}\,+b_ {out})+X, \tag{6}\]

where \(Y\) has the same size as \(X\), \(W_{\text{in}}\,\) and \(W_{\text{out}}\,\) are the weights. The activation function used is \(\mathrm{GELU}\)[20]. \(b_{in}\) and \(b_{out}\,\) are biases. In the implementation, all \(InputW+b\) are achieved by using a convolutional layer following a batch normalization operation. \(\mathrm{GraphConv}\) means aggregating and updating the discussed graph-level processing. The Grapher module is with a shortcut structure. The \(\mathrm{Droppath}\)[21] operation is used.

### Feed-forward Network

Feed-forward Networks (FFNs) are used to help with the feature transformation capacity and relief the over-smoothing phenomenon after the Grapher module. The FFN can be represented as

\[Z=\mathrm{Droppath}(\mathrm{GELU}\left(YW_{1}+b_{1}\right)W_{2}+b_{2})+Y, \tag{7}\]

where \(W_{1}\) and \(W_{2}\) are weights, \(b_{1}\) and \(b_{2}\) are biases. In the implementation, each feed-forward operation \(InputW+b\) is achieved by using a convolutional layer following a batch normalization operation. The \(\mathrm{Droppath}\) operation is used. The workflow of Grapher and FFN modules is shown in Figure2.

## 3 Experiments

### Datasets

**ISIC 2016** is a dataset of dermoscopic images of skin lesions. We used the dataset of lesion segmentation in this paper. There are 900 pairs of images and corresponding masks in the training set. In the testing set, there are 379 pairs. **ISIC 2017** is a dataset of dermoscopic images of skin lesions. We used the dataset of the lesion segmentation task, which contains images and corresponding masks. There are 2000 pairs in the training set, 150 in the validation set, and 600 in the testing set. The **Kvasir-SEG** dataset contains 1000 pairs of polyp images and masks. We split the dataset into training and testing sets with a ratio of 0.2 with a random state of 41.

### Implementation Details

The experiments are all done on PG500-216(V-100) with 32 GB memory. The training and validation set of ISIC 2016 and

Figure 1: The architecture of ViG-UNet: the basic modules are the stem block for visual embedding, Grapher Modules, Feed-forward Networks (FFNs) modules, downsampling modules in the encoder and upsampling modules in the decoder.

Figure 2: The workflow of Grapher and FFN modules: graph processing and feature transformation are applied
Kvair-SEG are split with a ratio of 0.2 with the random state 41. The total training epochs are 200 and the batch size is 4. The input images are all resized to 512\(\times\)512. The optimizer used is ADAM [10]. The initial learning rate is 0.0001 and a CosineAnnealingLR [11] scheduler is used. The minimum learning rate is 0.00001. Only rotation by 90 degrees clockwise for random times, flipping and normalization methods are used for augmentation. The evaluation metrics in validation are \(IOU\) of the lesions. A mixed loss combining binary cross entropy (BCE) loss and dice loss [1] is used in the training process:

\[\mathcal{L}=0.5BCE(\hat{y},y)+Dice(\hat{y},y)\]

We implemented ViG-UNet and six other U-Net variants for comparison experiments. The pre-trained Swin-T model of \(224\times 224\) input size on ImageNet 1k is used for Swin-UNet, and the pre-trained ViT-B/16 on ImageNet 21k is used for Trans-UNet.

### Results

The performances of different methods are shown in Table 2 and Table 3 with metrics of \(IoU\) and \(Dice\). The example segmentation results of different methods are displayed in Figure 3. From these experiments, we can see that on all three datasets, our approach performs best. And we can expect that if we use pre-trained models of ViG on ImageNet, the performance may be better. For the Swin-UNet, it's strange, but [1] also reports its low performance on small datasets. In conclusion, our method shows competitive performance compared to classical and state-of-the-art techniques.

We also calculate the parameters by using fvcore Python package with (1, 3, 512, 512) input. Admittedly, our model is larger than others and needs more computational resources.

## 4 Conclusion

In this work, we propose a ViG-UNet for 2D medical image segmentation, which has a GNN-based U-shaped architecture consisting of the encoder, the bottleneck, and the decoder with skip connections. Experiments are done on the ISIC 2016, the ISIC 2017 and the Kvasir-SEG dataset, whose results show

\begin{table}
\begin{tabular}{c c} \hline Methods & Parameters \\ \hline UNet & 7.8M \\ Attention-UNet & 8.7M \\ UNet++ & 9.2M \\ Trans-UNet & 92.3M \\ Swin-UNet & 27.3M \\ UNext & 1.5M \\ \hline ViG-UNet & 0.7G \\ \hline \end{tabular}
\end{table}
Table 4: Comparison of Parameters of Different Models

\begin{table}
\begin{tabular}{c c c c} \hline Methods & ISIC 2016 & ISIC 2017 & Kvasir SEG \\ \hline UNet & 0.8984 & 0.7708 & 0.8023 \\ Attention-UNet & 0.9058 & 0.7739 & 0.8065 \\ UNet++ & 0.9070 & 0.7768 & 0.8033 \\ Trans-UNet & 0.9158 & 0.8244 & 0.6439 \\ Swin-UNet & 0.8568 & 0.7914 & 0.4974 \\ UNext & 0.9103 & 0.8241 & 0.8122 \\ \hline ViG-UNet & **0.9206** & **0.8292** & **0.8188** \\ \hline \end{tabular}
\end{table}
Table 3: Comparison Experimental Results on ISIC 2016, ISIC 2017 and Kvasir SEG (using the \(Dice\) metric)

Figure 3: The example segmentation experimental results of different methods on ISIC 2016, ISIC 2017 and Kvasir-SEG datasets
that our method is effective.

## 5 Acknowledgements

This work was supported by a Grant from The National Natural Science Foundation of China (No. U21A20484).
