Snapshot Ensemble on Brain MRI Segmentation

Ensemble learning, where multiple models are trained and used to make a final prediction, is one of many ways to improve the performance of deep neural nets. While there are different approaches in generating ensemble models (e.g., by changing training data, pre-processing and post-processing methods, hyper parameters of a network or the network architecture itself), the majority of them are compute-intensive with training time increasing linearly with the number of models used in the ensemble. We investigate an unconventional ensemble approach to reduce training time with little compromise on performance. In this approach, multiple models are created by taking snapshots at different epochs while training a single model (hereinafter snapshot ensemble). The performance of this approach was evaluated using BraTS challenge 2021 dataset for multimodal brain tumor segmentation and compared to two other standard ensemble approaches where models were created using 20%-20% (Non-overlap) and 80%-20% (Overlap) train-test data splits. Experiments were carried out using two different deep neural nets (SegResNet and UNet) from MONAI library. Experiments were repeated five times using different random seeds, and mean (±standard deviation) dice score was reported as a performance metric. The mean dice scores for three ensemble approaches were: 0.770 ± 0.004 (Snapshot), 0.738 ± 0.003 (Non-Overlap) and 0.771 ± 0.003 (Overlap) for SegResNet, and 0.739 ± 0.003, 0.701 ± 0.007 and 0.743 ± 0.002 respectively for UNet. Our results show that SegResNet models were significantly better than UNet models (p < 0.01). In the SegResNet models, the dice scores of the snapshot and overlap ensembles were significantly higher than the dice score of non-overlap ensemble (p < 0.01), and there was no significant difference between the dice scores of snapshot and overlap ensembles. These results show that compared to overlap ensemble, the snapshot ensemble reduced training times by a factor of five without compromising performance, and can be used to obtain more sustainable models during medical imaging segmentation tasks.