High Resolution Medical Image Analysis with Spatial Partitioning

TL;DR

This paper introduces a spatial partitioning technique for training high-resolution 3D medical images with neural networks, enabling end-to-end processing of images up to 512x512x512 resolution without information loss.

Contribution

The authors develop a spatial partitioning method within Mesh-TensorFlow that distributes convolutional layer computations across multiple GPUs/TPUs, allowing high-resolution image analysis.

Findings

Successfully trained a 3D U-Net on 512x512x512 images.

First end-to-end approach for such high-resolution medical images.

Achieved efficient training without cropping or down-sampling.

Abstract

Medical images such as 3D computerized tomography (CT) scans and pathology images, have hundreds of millions or billions of voxels/pixels. It is infeasible to train CNN models directly on such high resolution images, because neural activations of a single image do not fit in the memory of a single GPU/TPU, and naive data and model parallelism approaches do not work. Existing image analysis approaches alleviate this problem by cropping or down-sampling input images, which leads to complicated implementation and sub-optimal performance due to information loss. In this paper, we implement spatial partitioning, which internally distributes the input and output of convolutional layers across GPUs/TPUs. Our implementation is based on the Mesh-TensorFlow framework and the computation distribution is transparent to end users. With this technique, we train a 3D Unet on up to 512 by 512 by 512 resolution data. To the best of our knowledge, this is the first work for handling such high resolution images end-to-end.