3D U2-Net: A 3D Universal U-Net for Multi-domain Medical Image Segmentation

Abstract

Fully convolutional neural networks like U-Net have been the state-of-the-art methods in medical image segmentation. Practically, a network is highly specialized and trained separately for each segmenta- tion task. Instead of a collection of multiple models, it is highly desirable to learn a universal data representation for different tasks, ideally a sin- gle model with the addition of a minimal number of parameters steered to each task. Inspired by the recent success of multi-domain learning in image classification, for the first time we explore a promising univer- sal architecture that handles multiple medical segmentation tasks and is extendable for new tasks, regardless of different organs and imaging modalities. Our 3D Universal U-Net (3D U2-Net) is built upon sepa- rable convolution, assuming that images from different domains have domain-specific spatial correlations which can be probed with channel- wise convolution while also share cross-channel correlations which can be modeled with pointwise convolution. We evaluate the 3D U2-Net on five organ segmentation datasets. Experimental results show that this uni- versal network is capable of competing with traditional models in terms of segmentation accuracy, while requiring only about 1% of the param- eters. Additionally, we observe that the architecture can be easily and effectively adapted to a new domain without sacrificing performance in the domains used to learn the shared parameterization of the universal network. We put the code of 3D U2-Net into public domain (https:// github.com/huangmozhilv/u2net