SAU-Net：基于U-Net和自注意力机制的医学图像分割方法

doi:10.12263/DZXB.20200984

摘要/Abstract

摘要：

基于深度学习的生物医学图像分割由于其精度的提高，可以更好地辅助医生做精确的诊断.目前主流的基于U-Net的分割模型通过多层卷积进行局部特征的提取，缺失了全局信息，使分割过于局部化而产生误差.本文通过自注意力机制和分解卷积策略对U-Net模型进行改进，提出一种新的深度分割网络SAU-Net，使用自注意力模块增加全局信息，将原U-Net中的级联结构改为逐像素相加，减小维度，降低计算量；提出一种快速简洁的分解卷积方法，将传统卷积分解为两路一维卷积，并加入残差连接强化上下文信息.在BRATS和Kaggle两个脑肿瘤数据集上进行的实验结果表明，SAU-Net在参数量和Dice系数上都有更优的性能.

关键词: 自注意力, 分解卷积, 医学图像分割, 深度学习, U-Net

Abstract:

Biomedical image segmentation based on deep learning can better help doctors make an accurate diagnosis due to its enhanced accuracy. At present, the U-Net-based mainstream segmentation model extracts local features through multi-layer convolutions, which lacks global information and leads to over-localized results with errors. This paper improves the U-Net model through the self-attention mechanism and decomposition convolution and proposes a new deep segmentation network called SAU-Net. The model uses the self-attention module to increase global information, and changes the cascade structure in the original U-Net to pixel-by-pixel addition in order to reduce the dimension and cut down the calculation cost. A fast and concise decomposition convolution method is proposed which integrates the traditional convolution into a two-way one-dimensional convolution, and the residual connection is added to enhance the context information. The experimental results conducted on the two brain tumor datasets of BRATS and Kaggle show that SAU-Net has better performance in terms of parameters and the Dice coefficients.

Key words: self-attention, decomposition convolution, medical image segmentation, deep learning, U-Net

中图分类号:

TP301.6

张淑军, 彭中, 李辉. SAU-Net：基于U-Net和自注意力机制的医学图像分割方法[J]. 电子学报, 2022, 50(10): 2433-2442.

ZHANG Shu-jun, PENG Zhong, LI Hui. SAU-Net: Medical Image Segmentation Method Based on U-Net and Self-Attention[J]. Acta Electronica Sinica, 2022, 50(10): 2433-2442.

图/表 15

参考文献 26

1	ALTAF F, ISLAM S M S, AKHTAR N, et al. Going deep in medical image analysis: Concepts, methods, challenges, and future directions[J]. IEEE Access, 2019, 7: 99540-99572.
2	RONNEBERGER O, FISCHER P, BROX T. U-Net: Convolutional Networks for Biomedical Image Segmentation[C]//Lecture Notes in Computer Science. Cham: Springer International Publishing, 2015: 234-241.
3	ZHOU Z W, SIDDIQUEE M M R, TAJBAKHSH N, et al. UNet++: A nested U-net architecture for medical image segmentation[J]. Deep Learning in Medical Image Analysis and Multimodal Learning for Clinical Decision Support, 2018, 11045: 3-11.
4	KAMNITSAS K, LEDIG C, NEWCOMBE V F J, et al. Efficient multi-scale 3D CNN with fully connected CRF for accurate brain lesion segmentation[J]. Medical Image Analysis, 2017, 36: 61-78.
5	GHAFOORIAN M, KARSSEMEIJER N, HESKES T, et al. Non-uniform patch sampling with deep convolutional neural networks for white matter hyperintensity segmentation[C]//2016 IEEE 13th International Symposium on Biomedical Imaging. Prague: IEEE, 2016: 1414-1417.
6	HU J, SHEN L, ALBANIE S, et al. Squeeze-and-excitation networks[J]. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2020, 42(8): 2011-2023.
7	XIE Y P, ZHANG Z Z, SAPKOTA M, et al. Spatial clockwork recurrent neural network for muscle perimysium segmentation[C]//International Conference on Medical Image Computing and Computer-Assisted Intervention. Cham: Springer, 2016: 185-193.
8	CHEN J X, YANG L, ZHANG Y Z, et al. Combining fully convolutional and recurrent neural networks for 3D biomedical image segmentation[C]//NIPS'16: Proceedings of the 30th International Conference on Neural Information Processing Systems. 2016: 3044-3052.
9	TSENG K L, LIN Y L, HSU W, et al. Joint sequence learning and cross-modality convolution for 3D biomedical segmentation[C]//2017 IEEE Conference on Computer Vision and Pattern Recognition. Honolulu: IEEE, 2017: 3739-3746.
10	VASWANI A, SHAZEER N, PARMAR N, et al. Attention is all You need[C]//NIPS'17: Proceedings of the 31st International Conference on Neural Information Processing Systems. New York: ACM, 2017: 6000-6010.
11	WANG X L, GIRSHICK R, GUPTA A, et al. Non-local neural networks[C]//2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition. Salt Lake City: IEEE, 2018: 7794-7803.
12	ZHANG H, GOODFELLOW I, METAXAS D, et al. Self-attention generative adversarial networks[C]//Proceedings of the 36th International Conference on Machine Learning. Long Beach: PMLR, 2019: 12744-12753.
13	HUANG Z L, WANG X G, HUANG L C, et al. CCNet: criss-cross attention for semantic segmentation[C]//2019 IEEE/CVF International Conference on Computer Vision(ICCV). Seoul: IEEE, 2019: 603-612.
14	FU J, LIU J, TIAN H J, et al. Dual attention network for scene segmentation[C]//2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition(CVPR). Long Beach: IEEE, 2019: 3141-3149.
15	ZHANG F, CHEN Y Q, LI Z H, et al. ACFNet: Attentional class feature network for semantic segmentation[C]//2019 IEEE/CVF International Conference on Computer Vision(ICCV). Seoul: IEEE, 2019: 6797-6806.
16	SINHA A, DOLZ J. Multi-scale self-guided attention for medical image segmentation[J]. IEEE Journal of Biomedical and Health Informatics, 2021, 25(1): 121-130.
17	OKTAY O, SCHLEMPER J, FOLGOC L L, et al. Attention U-Net: Learning where to look for the pancreas[C]//1st Conference on Medical Imaging with Deep Learning(MIDL 2018). Amsterdam: MIDL, 2018: 4861068.
18	LI R R, LI M M, LI J C, et al. Connection Sensitive Attention U-NET for Accurate Retinal Vessel Segmentation[EB/OL]. (2019)[2020]. .
19	QI K H, YANG H, LI C, et al. X-Net: Brain stroke lesion segmentation based on depthwise separable convolution and long-range dependencies[C]//International Conference on Medical Image Computing and Computer-Assisted Intervention. Cham: Springer, 2019: 247-255.
20	ZHANG H, ZHANG J W, ZHANG Q H, et al. RSANet: Recurrent slice-wise attention network for multiple sclerosis lesion segmentation[C]//International Conference on Medical Image Computing and Computer-Assisted Intervention. Cham: Springer, 2019: 411-419.
21	GAO H Y, WANG Z Y, CAI L, et al. ChannelNets: Compact and efficient convolutional neural networks via channel-wise convolutions[J]. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2021, 43(8): 2570-2581.
22	GAO H Y, YUAN H, WANG Z Y, et al. Pixel transposed convolutional networks[J]. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2020, 42(5): 1218-1227.
23	ROMERA E, ÁLVAREZ J M, BERGASA L M, et al. ERFNet: Efficient residual factorized ConvNet for real-time semantic segmentation[J]. IEEE Transactions on Intelligent Transportation Systems, 2018, 19(1): 263-272.
24	ZHANG X Y, ZHOU X Y, LIN M X, et al. ShuffleNet: an extremely efficient convolutional neural network for mobile devices[C]//2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition. Salt Lake City: IEEE, 2018: 6848-6856.
25	WANG Y, ZHOU Q, LIU J, et al. Lednet: A lightweight encoder-decoder network for real-time semantic segmentation[C]//2019 IEEE International Conference on Image Processing. Taipei: IEEE, 2019: 1860-1864.
26	MENZE B H, JAKAB A, BAUER S, et al. The multimodal brain tumor image segmentation benchmark(BRATS)[J]. IEEE Transactions on Medical Imaging, 2015, 34(10): 1993-2024.

数据量	方法	Train Dice loss	Train IoU	Test Dice loss	Test IoU
50例	U-Net	0.094 8	0.839 0	0.236 6	0.737 9
50例	SAU-Net	0.050 1	0.865 7	0.126 1	0.850 3
210例	U-Net	0.058 9	0.857 4	0.185 5	0.790 1
210例	SAU-Net	0.052 0	0.866 8	0.171 5	0.802 3

数据量	方法	Train Dice loss	Train IoU	Test Dice loss	Test IoU
50例	U-Net	0.094 8	0.839 0	0.236 6	0.737 9
50例	SAU-Net	0.050 1	0.865 7	0.126 1	0.850 3
210例	U-Net	0.058 9	0.857 4	0.185 5	0.790 1
210例	SAU-Net	0.052 0	0.866 8	0.171 5	0.802 3

分割方法	Train Dice Loss (均值±标准差)	Test Dice Loss (均值±标准差)
U-Net	0.396 11±0.000 23	0.441 89±0.000 25
SAU-Net-1	0.362 89±0.000 29	0.458 29±0.000 43
SAU-Net-FCB-1(Concat)	0.402 72±0.000 17	0.450 73±0.000 27
SAU-Net-FCB-1(Add)	0.406 21±0.000 20	0.427 71±0.000 30
SAU-Net-FCB-2(Concat)	0.362 81±0.000 32	0.469 14±0.000 40
SAU-Net-FCB-2(Add)	0.370 14±0.000 20	0.467 17±0.000 39

分割方法	Train Dice Loss (均值±标准差)	Test Dice Loss (均值±标准差)
U-Net	0.396 11±0.000 23	0.441 89±0.000 25
SAU-Net-1	0.362 89±0.000 29	0.458 29±0.000 43
SAU-Net-FCB-1(Concat)	0.402 72±0.000 17	0.450 73±0.000 27
SAU-Net-FCB-1(Add)	0.406 21±0.000 20	0.427 71±0.000 30
SAU-Net-FCB-2(Concat)	0.362 81±0.000 32	0.469 14±0.000 40
SAU-Net-FCB-2(Add)	0.370 14±0.000 20	0.467 17±0.000 39

模型	参数量
U-Net	39, 401, 026
SAU-Net	40, 699, 523
SAU-Net-FCB-1(Concat)	11, 087, 555
SAU-Net-FCB-1(Add)	8, 273, 091
SAU-Net-FCB-2(Concat)	11, 169, 796
SAU-Net-FCB-2(Add)	8, 355, 332