Progressive Growing of Patch Size: Resource-Efficient Curriculum Learning for Dense Prediction Tasks

TL;DR

This work introduces Progressive Growing of Patch Size (PGPS), a curriculum learning strategy that gradually increases the input patch size during training to convert easier segmentation tasks into harder ones. Implemented within nnU-Net and evaluated on the Medical Segmentation Decathlon, PGPS achieves substantial reductions in training runtime and CO2 emissions while maintaining or improving Dice scores on several tasks. The results show faster convergence and efficiency gains, with PGPS+ further leveraging larger batch sizes to boost performance on additional tasks. Overall, PGPS offers a practical, scalable approach to resource-efficient dense prediction in medical imaging, with potential applicability to other architectures and tasks.

Abstract

In this work, we introduce Progressive Growing of Patch Size, a resource-efficient implicit curriculum learning approach for dense prediction tasks. Our curriculum approach is defined by growing the patch size during model training, which gradually increases the task's difficulty. We integrated our curriculum into the nnU-Net framework and evaluated the methodology on all 10 tasks of the Medical Segmentation Decathlon. With our approach, we are able to substantially reduce runtime, computational costs, and CO2 emissions of network training compared to classical constant patch size training. In our experiments, the curriculum approach resulted in improved convergence. We are able to outperform standard nnU-Net training, which is trained with constant patch size, in terms of Dice Score on 7 out of 10 MSD tasks while only spending roughly 50% of the original training runtime. To the best of our knowledge, our Progressive Growing of Patch Size is the first successful employment of a sample-length curriculum in the form of patch size in the field of computer vision. Our code is publicly available at https://github.com/compai-lab/2024-miccai-fischer.

Figures (3)

• Figure 1: Proposed Progressive Growing of Patch Size curriculum. Our curriculum is illustrated for the lung cancer segmentation with cancer bounding boxes in yellow. A fully convolutional network is able to handle inputs of different sizes. Training with our curriculum starts by training the network with minimal patch size and, with training progress, increasing the patch size until the final patch size is reached. The ratio between foreground and background voxels is bigger for small patch sizes and decreases for bigger patch sizes. In contrast, the global context that can be inferred from the patch is growing with the patch size. For inference, the maximal patch size is used.
• Figure 2: Segmentation performance of models trained with Progressive Growing of Patch Size (PGPS), Random Patch Size Sampling (RPSS) and Constant Patch Size (CPS) for different numbers of training iterations per epoch on the MSD Lung Cancer Task. Models were trained with 10%, 25%, 50%, and 100% of the training iterations per epoch, while 100% represents the default 250 iterations per epoch of nnU-Net. Dice Scores are averaged over a 5-fold cross-validation. All three plots are referring to the same models but are evaluated regarding different measures.
• Figure 3: Performance of Progressive Growing of Patch Size (PGPS), PGPS with increased batch size (PGPS+) and Constant Patch Size (CPS) on the MSD tasks. Here, we show the fold-wise mean difference in Dice Score and its standard error of the 5-fold cross-validation. A positive difference represents wins of the PGPS curricula over the CPS training (here marked in dark blue). On this measure, PGPS outperforms CPS in 6 of 10 tasks, and PGPS+ outperforms CPS in 8 of 10 tasks.