ViFUNet: a Vision Flash based UNet for lung nodules segmentation task

Taoyu Chen, Fei Luo, Fu Zhou, Feng Liu, Yunfei Zha

Abstract: As an efficient computer-aided diagnosis technology, pulmonary nodule image segmentation plays an important role in greatly improving the identification efficiency. Due to their different shapes and sizes, there is always tremendous difficulty in the segmentation of lung nodules. The attention and Vision Transformer mechanisms are used in such scenarios, such as TransUNet, to solve the problem faced by traditional CNNs to obtain local information as well as have a good segmentation result. However, the addition to a network leads to it being bloated, increases the number of network parameters, slows down the training speed as well as makes it more prone to overfitting. In this paper, we propose an image segmentation network ViFUNet which is based on the Flash attention mechanism. Compared with the traditional Transformer, Flash adopts the fusion attention mechanism, which greatly reduces the number of network parameters and improves the network training speed. Experiments show that the convergence speed of ViFUNet is faster than that of the ViT-based network. In terms of segmentation accuracy, the Dice coefficient of ViFUNet’s lung nodule segmentation reaches 86.03%,which is 1.49% higher than U-Net and 2.66% higher than TransUNet, and it also performs well for different types of lung nodules. The average DSC score of pulmonary nodules less than 4mm in diameter reaches 85.20%, which is 1.37% higher than Atten-UNet, 5.03% higher than TransUNet, and 4.31% higher than U-Net. In addition, its parameter size is only 12. 7M, which is about half of the UNet and one-tenth of the TransUNet.