Search

Deep learning has achieved overwhelmingly better accuracy than existing methods in various fields. To further improve deep learning, a deeper and larger neural network (NN) model is indispensable, but graphics processing unit (GPU) memory is not large enough to train such NN models. One promising method to reduce GPU memory consumption is data swapping, which eases the burden on GPU memory by swapping out intermediate data from GPU memory to central processing unit (CPU) memory while the data are not necessary. However, this method introduces communication overhead in transferring data between GPU and CPU memory. Another method is recomputation, which discards intermediate data once and then computes them again when necessary. Unlike the data-swapping method, the re-computation method introduces additional computation but does not require CPU-GPU communication. Therefore, it may reduce CPU-GPU communication by introducing it effectively in the data-swapping method. In this paper, we developed a faster training method for large NN models by combining the re-computation and the data-swapping methods. This method edits the graph of TensorFlow automatically. We developed heuristics about which part of the graph should be recomputed. The heuristics divides the graph into sub-graphs and applies the re-computation in the decreasing order of the amount of swapping data in the sub-graphs. This heuristics is intended to reduce maximal communications by re-computing minimal sub-graphs. Our hybrid method improves performance by up to 15.5% in image size 8000×8000 of ResNet50, 15.3% in image size 7500×7500 of ResNet152, and 12.4% in image size 1000×1000 of DeepLabV3+ compared with the existing data-swapping method.

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.