3D U-Net Brain Tumor Segmentation Using VAE Skip Connection

Abstract

In clinical practice, the determination of the location, shape, and size of brain tumor can greatly assist the diagnosis, monitoring, and treatment of brain tumor. Therefore, accurate and reliable automatic brain tumor segmentation algorithm is of great significance for clinical diagnosis and treatment. With the rapid development of deep learning technology, more and more efficient image segmentation algorithms have also been applied in this field. It has been proven that U-Net model combined with variational auto-encoder can help to effectively regularize the shared encoder and thereby improve the performance. Based on the VAE-U-Net model, this paper proposes a structure called VAE skip connection. By fusing the position information contained in VAE branch into U-Net decoding stage, the network can retain more high-resolution detail information. In addition, we integrate ShakeDrop regularization into the networks to further alleviate the overfitting problem. The experimental results show that the networks after adding VAE skip connection and ShakeDrop can achieve competitive results on the BraTS 2018 dataset.