This is PyTorch implementation of Deep Residual 3D U-Net for MRI segmentation.
This is also the source code for Automatic post-stroke lesion segmentation on MR images using 3D residual convolutional neural network
- We developed a deep learning model for MRI 3D lesion segmentation of chronic stroke.
- Our novel zoom-in&out strategy increases performance and accelerates optimization.
- High performance was achieved in both volumetric and surface-based metrics.
We used a two-stage zoom-in&out strategy to first train them on small volumes, and then we finetuned the models on larger volumes. Advantages are:
- Training models with smaller volumes can have a regularizing effect due to performing data augmentation by randomly extracting diverse sub-volumes from original volumes.
- A “zoom-in” step is a computationally inexpensive option and can utilize sub-optimal graphic processing units (GPUs) for the task. By feeding smaller volumes to older but more accessible GPUs, models can be trained in parallel, and, as a result, are faster.
- The “zoom-out” stage involves showing models larger volumes to learn from the broader context of input images and improves the robustness of the model.
| Optimization Stage | Zoom-In Stage | Zoom-Out Stage |
|---|---|---|
| Input volume size (mm3) | 128 × 128 × 128 (24% sub-volume) | 144 × 172 × 168 (48% sub-volume) |
| Training Length | 1200 epochs | 150 epochs |
| Initial learning rate | 1.00E−03 | 1.00E−04 |
| Optimizer | Adam optimizer and cosine annealing with warm restart scheduler | |
| GPU | Nvidia Titan Xp with 12 GB memory | Nvidia Titan RTX with 24 GB memory |
| Methods | microDSC | DSC | HD (mm) | ASSD (mm) | TPR | Precision |
|---|---|---|---|---|---|---|
| 3D-ResU-Net w/o zoom-out | 0.73 | 0.60 (0.47–0.73) | 35.1 (20.4–51.3) | 7.6 (3.7–12.3) | 0.83 (0.71–0.91) | 0.54 (0.39–0.67) |
| 3D-ResU-Net w/ zoom-in&out | 0.79 | 0.64 (0.51–0.76) | 20.4 (10.0–33.3) | 3.6 (1.7–6.2) | 0.81 (0.68–0.89) | 0.62 (0.48–0.74) |
| Δ | +0.06 | +0.04 | −14.7 | −4.0 | −0.02 | +0.08 |
- Python 3.6
- PyTorch 1.3
- NiBabel
- scikit-image
- scikit-learn
python train.py --data_dir <path_to_dataset> --save_dir <path_to_snapshots>
Zoom-in stage
python train.py --data_dir <path_to_dataset> --save_dir <path_to_snapshots> -v xi yi ziwhere xi, yi, zi are the size of input volume in the zoom-in stage.
Zoom-out stage
python train.py --data_dir <path_to_dataset> --save_dir <path_to_snapshots> -v xo yo zo --resume_model <path_to_a_snapshot>where xo, yo, zo are the size of input volume in the zoom-out stage.
Here are our suggestions for selecting an optimal volume size in Zoom-In&Out training strategy
- Visualizing typical images from the dataset, and then reviewing the spatial distribution of target lesions to estimate the possible effective input volume size for a specific application. (In our study, we observed in early experiments that using volume sizes significantly smaller than those in our final configuration does not achieve a competitive performance. This performance gap may indicate that volumes cropped for training during these experiments were overwhelmed by regions without positive stroke lesions, which could have negatively biased our model optimization in the early experiments.)
- Performing a systematic search on proportions of original volume size (e.g., 25%, 50%, 75% ) to find an optimal configuration.
python test.py --data_dir <path_to_dataset> --model_path <path_to_a_snapshot>
3dMRISegmentation is an open-source library and is licensed under the GNU General Public License (v3).
If you are using this library, please cite:
Naofumi Tomita, Steven Jiang, Matthew E. Maeder, and Saeed Hassanpour. "Automatic Post-Stroke Lesion Segmentation on MR Images using 3D Residual Convolutional Neural Network." NeuroImage: Clinical (2020): 102276.