High Resolution Medical Image Segmentation Using Data-Swapping Method

Abstract

Deep neural network models used for medical image segmentation are large because they are trained with high-resolution three-dimensional (3D) images. Graphics processing units (GPUs) are widely used to accelerate training. However, the memory on a GPU is not large enough to train the models. A popular approach to tackling this problem is the patch-based method, which divides a large image into small patches and trains the models with these small patches. However, this method degrades the segmentation quality if a target object spans multiple patches. In this paper, we propose a novel approach for 3D medical image segmentation that utilizes the data-swapping method, which swaps out intermediate data from GPU memory to CPU memory to enlarge the effective GPU memory size for training high-resolution 3D medical images without patching. We enhanced the existing data-swapping method by introducing swapping inside forward propagation and selective swapping of analysis path in order to train 3D U-Net effectively. We applied this approach to train 3D U-Net with full-size images of 192 × 192 × 192 voxels for a brain tumor dataset. Compared with the patch-based method for patches of 128 × 128 × 128 voxels, our approach improved the mean Dice score by 3.9 percentage points and 4.1 percentage points when detecting a whole tumor sub-region and a tumor core sub-region, respectively. The total training time was reduced from 164 h to 47 h, resulting in an acceleration of 3.53 times.