The following sections provide greater details of the dataset, running training and inference, and the training results.

Scripts and sample code

In the root directory, the most important files are:

• main.py: Serves as the entry point to the application. Encapsulates the training routine.
• Dockerfile: Container with the basic set of dependencies to run U-Net.
• requirements.txt: Set of extra requirements for running U-Net.
• preprocess_data.py: Converts the dataset to tfrecord format for training.

The dataset/ folder contains the necessary tools to train and perform inference using U-Net. Its main components are:

• data_loader.py: Implements the data loading and augmentation.
• transforms.py: Implements the data augmentation functions.
• preprocess_data.py: Implements the data conversion and pre-processing functionality.

The runtime/ folder contains scripts with training and inference logic. Its contents are:

• arguments.py: Implements the command-line arguments parsing.
• hooks.py: Collects different metrics to be used for benchmarking and testing.
• parse_results.py: Defines a set of functions used for parsing the partial results.
• setup.py: Defines a set of functions to set the environment up.

The model/ folder contains information about the building blocks of 3D-UNet and the way they are assembled. Its contents are:

• layers.py: Defines the different blocks that are used to assemble 3D-UNet.
• losses.py: Defines the different losses used during training and evaluation.
• model_fn.py: Defines the computational graph to optimize.
• unet3d.py: Defines the model architecture using the blocks from the layers.py file.

Other folders included in the root directory are:

• scripts/: Provides examples for training and benchmarking U-Net
• images/: Contains the model diagram

Parameters

The complete list of the available parameters for the main.py script contains:

• --exec_mode: Select the execution mode to run the model (default: train). Modes available:
  • train - trains a model and stores checkpoints in the directory passed using --model_dir
  • evaluate - loads checkpoint (if available) and performs evaluation on validation subset (requires --fold other than None).
  • train_and_evaluate - trains model from scratch and performs validation at the end (requires --fold other than None).
  • predict - loads checkpoint (if available) and runs inference on the test set. Stores the results in the --model_dir directory.
  • train_and_predict - trains model from scratch and performs inference.
• --model_dir: Set the output directory for information related to the model (default: /results).
• --log_dir: Set the output directory for logs (default: None).
• --data_dir: Set the input directory containing the dataset (default: None).
• --batch_size: Size of each minibatch per GPU (default: 1).
• --fold: Selected fold for cross-validation (default: None).
• --max_steps: Maximum number of steps (batches) for training (default: 16000).
• --seed: Set random seed for reproducibility (default: 0).
• --log_every: Log performance every n steps (default: 100).
• --learning_rate: Model's learning rate (default: 0.0002).
• --augment: Enable data augmentation (disabled by default).
• --benchmark: Enable performance benchmarking (disabled by default). If the flag is set, the script runs in a benchmark mode - each iteration is timed and the performance result (in images per second) is printed at the end. Works for both train and predict execution modes.
• --warmup_steps: Used during benchmarking - the number of steps to skip (default: 200). First iterations are usually much slower since the graph is being constructed. Skipping the initial iterations is required for a fair performance assessment.
• --resume_training: Whether to resume training from a checkpoint, if there is one (disabled by default)
• --xla: Enable accelerated linear algebra optimization (disabled by default).
• --amp: Enable automatic mixed precision (disabled by default).

Command line options

To see the full list of available options and their descriptions, use the -h or --help command-line option, for example:

python main.py --help

The following example output is printed when running the model:

usage: main.py [-h] --data_dir DATA_DIR --model_dir MODEL_DIR
               [--exec_mode {train,evaluate,train_and_evaluate,predict}]
               [--benchmark] [--max_steps MAX_STEPS]
               [--learning_rate LEARNING_RATE] [--log_every LOG_EVERY]
               [--log_dir LOG_DIR] [--loss {dice,ce,dice+ce}]
               [--warmup_steps WARMUP_STEPS][--resume_training] 
               [--augment] [--batch_size BATCH_SIZE] [--fold FOLD] 
               [--amp] [--xla]
 
UNet-3D
 
optional arguments:
 -h, --help            show this help message and exit
 --model_dir MODEL_DIR
                       Output directory for information related to the model
 --data_dir DATA_DIR   Input directory containing the dataset for training
                       the model
 --exec_mode {train,train_and_predict,predict,evaluate,train_and_evaluate}
                       Execution mode of running the model
 --log_dir LOG_DIR     Output directory for training logs
 --batch_size BATCH_SIZE
                       Size of each minibatch per GPU
 --learning_rate LEARNING_RATE
                       Learning rate coefficient for AdamOptimizer
 --fold                Fold number
                       Chosen fold for cross-validation. Use None to disable
                       cross-validation
 --max_steps MAX_STEPS
                       Maximum number of steps (batches) used for training
 --log_every LOG_EVERY
                       Log performance every n steps
 --warmup_steps WARMUP_STEPS
                       Number of warmup steps
 --resume_training     Whether to resume training from the checkpoint
 --seed SEED           Random seed
 --augment             Perform data augmentation during training
 --benchmark           Collect performance metrics during training
 --amp                 Train using TF-AMP
 --xla                 Train using XLA

Getting the data

The 3D-UNet model was trained in the Brain Tumor Segmentation 2019 dataset. Test images provided by the organization were used to produce the resulting masks for submission. Upon registration, the challenge's data is made available through the https//ipp.cbica.upenn.edu service.

The dataset consists of 335 240x240x155 nifti volumes. Each volume is represented by 4 modalities and a corresponding segmentation mask. The modalities are:

• Native T1-weighted (T1),
• Post-contrast T1-weighted (T1Gd),
• Native T2-weighted (T2),
• T2 Fluid Attenuated Inversion Recovery (FLAIR).

Each voxel in a segmentation mask belongs to one of four classes:

• 0 corresponds to healthy tissue or background,
• 1 indicates the presence of the necrotic and non-enhancing tumor core (TC),
• 2 indicates the presence of the peritumoral edema (ED),
• 4 indicates the presence of the GD-enhancing tumor (ET).

The objective is to produce a set of masks that segment the data as accurately as possible. The results are expected to be submitted as a 12-bit nifti 3D image, with values corresponding to the underlying class.

Dataset guidelines

The training and test datasets are given as 3D nifti volumes that can be read using the Nibabel library and NumPy (both packages are installed by the Dockerfile).

Initially, all modalities are loaded, stacked and converted into 240x240x155x4 NumPy arrays using Nibabel. To decrease the size of the dataset, each volume is clipped to 85% of the maximal value, normalized to 255 for each modality separately, casted to 8-bit, grouped by 4 volumes, and saved as a tfrecord file. The process of converting from nifti to tfrecord can be found in the preprocess_data.py script.

The tfrecord files are fed to the model through tf.data.TFRecordDataset() to achieve high performance.

The foreground voxel intensities then z-score normalized, whereas labels are one-hot encoded for their later use in dice or pixel-wise cross-entropy loss, becoming 240x240x155x4 tensors.

If augmentation is enabled, the following set of augmentation techniques are applied:

• Random horizontal flipping
• Random 128x128x128x4 crop
• Random brightness shifting

In addition, random vertical flip and random gamma correction augmentations were implemented, but are not used. The process of loading, normalizing and augmenting the data contained in the dataset can be found in the data_loader.py script.

Multi-dataset

This implementation is tuned for the Brain Tumor Segmentation 2019 dataset. Using other datasets is possible, but might require changes to the code (data loader) and tuning some hyperparameters (e.g. learning rate, number of iterations).

In the current implementation, the data loader works with tfrecord files. It should work seamlessly with any dataset containing 3D data stored in tfrecord format, as long as features (with corresponding mean and standard deviation) and labels are stored as bytestream in the same file as X, Y, mean, and stdev. See the data pre-processing script for details. If your data is stored in a different format, you will have to modify the parsing function in the dataset/data_loader.py file. For a walk-through, check the TensorFlow tf.data API guide