AutoML Techniques for Medical Image Segmentation

In computer vision tasks such as image classification and segmentation convolutional neural networks are becoming a de-facto state of the art. More and more complex models are becoming available to solve those challenging tasks. As those complex neural networks are prone to overfitting they require large annotated datasets to be trained on. In medical imaging, such large annotated datasets are not common. This is especially true for segmentation where medical experts need to label each pixel in a time-consuming manner. A commonly used approach to tackle this shortage of data is the application of data augmentation. To keep the generalization abilities of the network data augmentation tries to sample additional data from the data distribution. The current state of the art to accomplish this is to manually design a sequence of image transformations which are randomly applied to the training data. This manual design requires domain knowledge and results in suboptimal choices, caused by its complexity and model dependence. Therefore the demand for an automatic, task-specific, and data-driven augmentation strategy arises. This allows for better generalization abilities of the network as well as important insight into the data itself. In this work, we propose an efficient optimization framework to automatically design an augmentation strategy based on the model and data directly. The augmentation strategy is represented as a non-commutative sequence of image transformations, defined as operator, probability, and magnitude tuples. The search strategy utilizes an autoencoder network to relax the discrete search problem to a continuous one. This is accomplished by training a value estimator to predict the performance of a sequence's latent space representation and utilizing gradient ascent. This performance prediction is based on a new measure that describes the generalization ability of the target network directly. We apply this automation to a challenging medical image segmentation problem and show the benefits for the network's performance.