clara_pt_fed_learning_brain_tumor_mri_segmentation

Model Overview

A federated learning demo for volumetric (3D) segmentation of brain tumors from multimodal MRIs based on BraTS 2018 data.

Note: The 4.1 version of this model is only compatible with the 4.1 version of the Clara Train SDK container

Federated Learning

Two clients share the global training models via a federated learning server for every epoch of local training.

The overall illustration of the federated training is shown below:

For more information, please see also: Clara Train example notebooks and Nvidia FLARE.

Model Architecture

The local model for each federated client utilized a similar approach described in 3D MRI brain tumor segmentation using autoencoder regularization, which was a winning method in BraTS2018 [1]. The client's training was performed with the following:

Training

The training was performed with the following:

• GPU: At least 16GB of GPU memory.
• Actual Model Input: 224 x 224 x 144
• AMP: True
• Optimizer: Adam
• Learning Rate: 1e-4
• Loss: DiceLoss

Dataset

The model is trained to segment 3 nested subregions of primary (gliomas) brain tumors: the "enhancing tumor" (ET), the "tumor core" (TC), the "whole tumor" (WT) based on 4 input MRI scans (T1c, T1, T2, FLAIR). The ET is described by areas that show hyperintensity in T1c when compared to T1, but also when compared to "healthy" white matter in T1c. The TC describes the bulk of the tumor, which is what is typically resected. The TC entails the ET, as well as the necrotic (fluid-filled) and the non-enhancing (solid) parts of the tumor. The WT describes the complete extent of the disease, as it entails the TC and the peritumoral edema (ED), which is typically depicted by hyperintense signal in FLAIR.

The dataset is available at "Multimodal Brain Tumor Segmentation Challenge (BraTS) 2018." The provided labelled data was partitioned, based our own split, into training (200 studies), validation (42 studies) and testing (43 studies) datasets.

For more detailed descriptions of tumor regions, please see the Multimodal Brain Tumor Segmentation Challenge (BraTS) 2018 data page at https://www.med.upenn.edu/sbia/brats2018/data.html.

Data Preparation

When deploying the MMAR to the clients (site-1 and site-2), the "DATASET_JSON" value in config/environment.json should be changed to "config/seg_brats18_datalist_client_a.json" for site-1, and "config/seg_brats18_datalist_client_b.json" for site-2.

Similarly, when evaluating the models, the commands/validate.sh is by default configured to run the MMAR model with "config/seg_brats18_datalist_client_a.json" testing data. The "DATASET_JSON" should be adjusted to the target data list if needed.

Performance

The model was trained with 144 High-grade for site-1 (validation 30; testing 36) and 56 Low-grade for site-2 (validation 12; testing 7), with our own split, as shown in the datalist json files in config folder.

Training:

A figure of two clients training with 600 federated rounds:

Validation:

A figure of local site validation of two clients during the 600-round federated learning:

Cross-site Validation:

The following report uses TC to denote "Tumor Core", WT to denote "Whole Tumor", ET to denote "Enhancing Tumor". The results are based on the testing sets (30 cases on site-1 and 12 cases on site-2).

Global model at the end of 600 federated rounds:

Site-1: TC: 0.8790 WT: 0.9180 ET: 0.8064
Site-2: TC: 0.6223 WT: 0.8841 ET: 0.5908
AVE:    TC: 0.8057 WT: 0.9083 ET: 0.7448

Global best model selected via IntimeModelSelectionHandler:

Site-1: TC: 0.8812 WT: 0.9152 ET: 0.7877
Site-2: TC: 0.6053 WT: 0.8683 ET: 0.5961
AVE:    TC: 0.8024 WT: 0.9018 ET: 0.7329

Site-1 best model selected via site-1 local validation:

Site-1: TC: 0.8811 WT: 0.9157 ET: 0.7876
Site-2: TC: 0.6427 WT: 0.8882 ET: 0.5109
AVE:    TC: 0.8130 WT: 0.9078 ET: 0.7085

Site-2 best model selected via site-2 local validation:

Site-1: TC: 0.8798 WT: 0.9137 ET: 0.7895
Site-2: TC: 0.6161 WT: 0.8676 ET: 0.5009
AVE:    TC: 0.8045 WT: 0.9005 ET: 0.7070

The cross-site average (AVE) is computed from (Site-1_mean x 30 + Site-2_mean x 12) / (30 + 12). An example of the cross-site validation results are available at docs/example_cross_site_eval/.

How to Use this Model

The model was validated with NVIDIA hardware and software. For hardware, the model can run on any NVIDIA GPU with memory greater than 16 GB. For software, this model is usable only as part of Transfer Learning & Annotation Tools in Clara Train SDK container. Find out more about Clara Train at the Clara Train Collections on NGC.

Full instructions for the training and validation workflow can be found in our documentation

Input

Input: 4 channel MRI (4 aligned MRIs T1c, T1, T2, FLAIR at 1x1x1 mm)

Preprocessing:

1. Normalizing to unit std with zero mean
2. Randomly cropping to 224 x 224 x 144

Augmentation in training:

1. Randomly spatial flipping
2. Randomly scaling and shifting intensity of the volume

Output

Output: 3 channels

• Label 0: TC tumor subregion
• Label 1: WT tumor subregion
• Label 2: ET tumor subregion

Limitations

This training and inference pipeline was developed by NVIDIA. It is based on a segmentation model developed by NVIDIA researchers. This research use only software has not been cleared or approved by FDA or any regulatory agency. Clara pre-trained models are for developmental purposes only and cannot be used directly for clinical procedures.

Reference

[1] Myronenko, Andriy. "3D MRI brain tumor segmentation using autoencoder regularization." International MICCAI Brainlesion Workshop. Springer, Cham, 2018. https://arxiv.org/abs/1810.11654.

License

End User License Agreement is included with the product. Licenses are also available along with the model application zip file. By pulling and using the Clara Train SDK container and downloading models, you accept the terms and conditions of these licenses.