Source-linked AI summary
Optimized U-Net for Brain Tumor Segmentation
Michał Futrega, Alexandre Milesi, Michal Marcinkiewicz, Pablo Ribalta
TL;DR
Automatic brain tumor segmentation is difficult because tumors vary across multimodal MRI and manual annotation is costly and inconsistent. The paper uses extensive ablations to optimize a U-Net architecture, training schedule, and post-processing strategy for BraTS21. U-Net with deep supervision and further architectural and input modifications achieved the best reported results and won the validation phase.
Problem
Automatic brain tumor segmentation remains difficult because tumors vary across multimodal MRI, while manual segmentation is time-consuming and inconsistent.
Method
The study performs extensive ablations of U-Net variants, training modifications, encoder depth, convolutional channels, input encoding, and post-processing on BraTS21.
Results
U-Net with deep supervision yields the best results and is further improved by an additional foreground input channel, increased encoder depth and convolutional channels, and post-processing.
Takeaways & Limitations
The optimized deeper U-Net with deep supervision, modified channels, and an additional foreground one-hot encoding channel won the BraTS21 validation phase.
Abstract
from arXiv · showhide
We propose an optimized U-Net architecture for a brain tumor segmentation task in the BraTS21 challenge. To find the optimal model architecture and the learning schedule, we have run an extensive ablation study to test: deep supervision loss, Focal loss, decoder attention, drop block, and residual connections. Additionally, we have searched for the optimal depth of the U-Net encoder, number of convolutional channels and post-processing strategy. Our method won the validation phase and took third place in the test phase. We have open-sourced the code to reproduce our BraTS21 submission at the NVIDIA Deep Learning Examples GitHub Repository.
1 Introduction
Automatic brain tumor segmentation is difficult because tumors vary in shape and appearance across multimodal MRI, while manual annotation is expertise-intensive, time-consuming, error-prone, and inconsistent. This motivates systematic selection of U-Net architectures and training schedules using the comparatively large BraTS21 dataset.
- Automatic brain tumor segmentation is challenging because gliomas vary in shape and appearance across multimodal MRI.
- Manual tumor segmentation requires substantial expertise, takes considerable time, and can suffer from human error, inconsistency, and poor reproducibility.
- U-Net-based encoder-decoder architectures perform strongly in brain tumor segmentation, but the many available variants make selecting an optimal architecture difficult.
- Training-schedule design is also critical for obtaining optimal performance from the selected neural network architecture.
- BraTS21 provides 2,040 examples across training, validation, and test sets, enabling performance comparisons among different U-Net variants.
- The study uses extensive ablations to select a U-Net variant and training schedule, testing architectural variants, losses, regularization, and model modifications.
2 Method
The method uses a 3D U-Net-based segmentation pipeline with preprocessing, augmentation, and systematic comparison of several encoder-decoder architectures. The selected designs include a basic U-Net, residual and attention variants, SegResNetVAE, and UNETR, with encoder depth and channel counts further optimized.
- Data preprocessing and augmentation: The BraTS21 volumes were stacked across four MRI modalities, cropped to remove redundant background, normalized within non-zero regions, and augmented during training.Augmentations included biased cropping, zooming, and flips; a foreground one-hot channel was also added to distinguish background from normalized voxels.
- Architecture comparison: The architecture study compared U-Net, Attention U-Net, Residual U-Net, SegResNetVAE, and UNETR as alternative encoder-decoder designs.Attention U-Net adds decoder attention gates, Residual U-Net adds residual connections, SegResNetVAE uses residual blocks and a training-only reconstruction branch, and UNETR replaces encoder convolutions with self-attention.
- U-Net architecture: The basic U-Net uses a contracting encoder that reduces spatial dimensions and an expanding decoder that restores them through upsampling and skip connections.Decoder feature maps are concatenated with encoder maps from corresponding spatial levels before further transformation.
- Deep supervision: Two additional decoder output heads support deep supervision by computing losses at multiple decoder levels.The architecture figure marks these auxiliary heads in green, while the method describes deep supervision as computing loss functions on different decoder levels.
- Architecture optimization: The selected basic U-Net was further improved by increasing encoder depth from 6 to 7 and changing the convolution channels to 64, 96, 128, 192, 256, 384, 512.The original six-level nnU-Net baseline used channels 32, 64, 128, 256, 320, 320.
3 Results
The experiments selected a baseline U-Net, then improved it with deep supervision, encoder modifications, foreground encoding, and post-processing. The resulting model won the BraTS21 validation phase.
- Architecture and evaluation: The experiments compared baseline U-Net, UNETR, and SegResNetVAE architectures, using five-fold cross-validation and checkpoint ensembling for evaluation.Training used Adam with warmup and cosine annealing, while inference averaged predictions from the two best checkpoints per fold.
- Architecture and evaluation: Baseline U-Net achieved the highest score among the compared architectures, while SegResNetVAE required three times longer training because of its VAE branch.The similar score of SegResNetVAE did not justify its additional training cost, so U-Net was selected for further exploration.
- Ablation studies: 0.9149 was the 5-fold average Dice score for deep supervision, compared with 0.9130 for baseline U-Net; other tested extensions did not significantly improve the baseline.The extensions included decoder attention, residual connections, drop block, and Focal+Dice loss.
- Encoder and input modifications: 0.9156 was achieved when deeper encoding, modified convolution channels, and a foreground one-hot input channel were applied together.Each modification separately slightly improved the deep-supervision baseline, while their combination produced the highest reported score in this experiment.
- Post-processing: The selected post-processing strategy reassigned small or globally sparse enhanced-tumor regions to NCR when their mean probability was below 0.9.It used thresholds of fewer than 16 voxels for connected components and fewer than 73 total enhanced-tumor voxels.
- Challenge result: The deeper U-Net with deep supervision, modified channels, and foreground one-hot encoding won the BraTS21 validation phase.Predictions on the validation dataset distinguish background, NCR, ED, and ET classes by color.
4 Conclusions
The study found that U-Net with deep supervision performed best among the tested variants and improved further with encoder, channel, input, and post-processing modifications.
- Conclusion: U-Net with deep supervision yielded the best results among the tested U-Net variants and training modifications.The experiments included basic U-Net, UNETR, SegResNetVAE, Residual U-Net, Attention U-Net, drop block, and Focal loss.
- Conclusion: Performance improved further by adding foreground one-hot encoding, increasing encoder depth and convolutional channels, and designing a post-processing strategy.