Model Overview

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

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 our Clara Train Federated Learning example notebooks.

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, the "DATASET_JSON" value in config/environment.json should be changed to "config/seg_brats18_datalist_client_a.json" for client A, and "config/seg_brats18_datalist_client_b.json" for client B.

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 client A (validation 30; testing 36) and 56 Low-grade for client B (validation 12; testing 7), with our own split, as shown in the datalist json files in config folder.

Training:

Validation:

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 Client A and 12 cases on client B).

Global Best Model

Data-A: TC: 0.8707 WT: 0.9146 ET: 0.7931
Data-B: TC: 0.5929 WT: 0.8793 ET: 0.5617
AVE: TC: 0.7913 WT: 0.9045 ET: 0.7270

Client-A Best Model

Data-A: TC: 0.8655 WT: 0.9157 ET: 0.7797
Data-B: TC: 0.5604 WT: 0.8514 ET: 0.5356
AVE: TC: 0.7783 WT: 0.8973 ET: 0.7100

Client-B Best Model

Data-A: TC: 0.8781 WT: 0.9170 ET: 0.7921
Data-B: TC: 0.5360 WT: 0.8469 ET: 0.4472
AVE: TC: 0.7804 WT: 0.8970 ET: 0.6936

The cross-site average (AVE) is computed from (client_A_mean x 30 + client_B_mean x 12) / (30 + 12) where client A has 30 test cases, client B has 12.

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.