Dynamic Data Augmentation via MCTS for Prostate MRI Segmentation

Xinyue Xu,Yuhan Hsi,Haonan Wang,Xiaomeng Li
DOI: https://doi.org/10.1007/978-981-99-8141-0_21
2023-08-31
Abstract:Medical image data are often limited due to the expensive acquisition and annotation process. Hence, training a deep-learning model with only raw data can easily lead to overfitting. One solution to this problem is to augment the raw data with various transformations, improving the model's ability to generalize to new data. However, manually configuring a generic augmentation combination and parameters for different datasets is non-trivial due to inconsistent acquisition approaches and data distributions. Therefore, automatic data augmentation is proposed to learn favorable augmentation strategies for different datasets while incurring large GPU overhead. To this end, we present a novel method, called Dynamic Data Augmentation (DDAug), which is efficient and has negligible computation cost. Our DDAug develops a hierarchical tree structure to represent various augmentations and utilizes an efficient Monte-Carlo tree searching algorithm to update, prune, and sample the tree. As a result, the augmentation pipeline can be optimized for each dataset automatically. Experiments on multiple Prostate MRI datasets show that our method outperforms the current state-of-the-art data augmentation strategies.
Image and Video Processing,Computer Vision and Pattern Recognition,Machine Learning
What problem does this paper attempt to address?
The paper aims to address the issue of data augmentation in prostate MRI segmentation. Specifically, the researchers propose a novel dynamic data augmentation method (Dynamic Data Augmentation, DDAug) to improve the generalization ability and segmentation performance of models in prostate MRI segmentation tasks. Traditionally, manually configuring data augmentation strategies is time-consuming and difficult to adapt to different datasets, while automatic data augmentation, although solving this problem, usually incurs significant GPU overhead. To address these issues, DDAug utilizes the Monte-Carlo Tree Search (MCTS) algorithm to construct a hierarchical tree structure representing various data augmentation operations, and optimizes the augmentation strategy for each dataset by updating, pruning, and sampling the tree structure. Experimental results show that DDAug outperforms the current state-of-the-art data augmentation strategies on multiple prostate MRI datasets and is comparable to the original nnUNet in terms of computational cost, demonstrating high efficiency. Additionally, compared to other methods, DDAug achieves better average Dice similarity coefficient with only a few operations, proving its effectiveness in enhancing the model's generalization ability.