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.
• Dockerfile: Container with the basic set of dependencies to run V-Net
• requirements.txt: Set of extra requirements for running V-Net
• download_data.py: Automatically downloads the dataset for training

The utils/ folder encapsulates the necessary tools to train and perform inference using V-Net. Its main components are:

• runner.py: Implements the logic for training and inference
• data_loader.py: Implements the data loading and augmentation
• hooks/profiling_hook.py: Collects different metrics to be used for benchmarking
• hooks/training_hook.py: Collects different metrics to be used for training
• hooks/evaluation_hook.py: Collects different metrics to be used for testing
• var_storage.py: Helper functions for TF-AMP

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

• layers.py: Defines the different blocks that are used to assemble V-Net
• vnet.py: Defines the model architecture using the blocks from the layers.py script

Other folders included in the root directory are:

• dllogger/: Contains the utilities for logging
• examples/: Provides examples for training and benchmarking V-Net
• images/: Contains a 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_and_predict)
• --data_normalization: Select the type of data normalization (default: zscore)
• --activation: Select the activation to be used in the network (default: prelu)
• --resize_interpolator: Select the interpolator for image resizing (default: nearest_neighbor)
• --loss: Loss function to be utilized for training (default: dice)
• --normalization_layer: Type of normalization layer to be used in the model (default: batchnorm)
• --pooling: Type of pooling layer to be used in the model (default: conv_pool)
• --upsampling: Type of upsampling layer to be used in the model (default: transposed_conv)
• --seed: Random seed value (default: 0)
• --input_shape: Target resize dimension for input samples (default: [32 32 32])
• --upscale_blocks: Number of upscale blocks with the depth of their residual component (default: [3 3 3])
• --downscale_blocks: Number of downscale blocks with the depth of their residual component (default: [3 3])
• --model_dir: Set the output directory for information related to the model (default: result/)
• --convolution_size: Size of the convolutional kernel filters (default: 3)
• --batch_size: Number of samples processed per execution step
• --log_every: Log every this number of epochs (default: 100)
• --warmup_steps: Initial number of steps that will not be benchmarked as the model starts running (default: 200)
• --train_epochs: Number of times that training will go through the entire dataset
• --optimizer: Optimizer to be used during training (default: rmsprop)
• --base_lr: Model's learning rate (default: 0.01)
• --momentum: Momentum coefficient for model's optimizer (default: 0.99)
• --train_split: Proportion of the dataset that will become the training set (default: 0.9)
• --split_seed: Random seed for the splitting of the dataset between training and validation
• --model_dir: Path where checkpoints and information related to the model will be stored
• --data_dir: Path to the dataset
• --augment: Enable data augmentation (default: False)
• --benchmark: Enable performance benchmarking (default: False)
• --amp: Enable automatic mixed precision (default: False)
• --xla: Enable xla (default: False)

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 
usage: main.py [-h]
               --exec_mode {train,predict,train_and_predict,train_and_evaluate}
               [--data_normalization {zscore}]
               [--activation {relu}]
               [--resize_interpolator {linear}]
               [--loss {dice}]
               [--normalization_layer {batchnorm}]
               [--pooling {conv_pool}]
               [--upsampling {transposed_conv}]
               [--seed SEED]
               [--input_shape INPUT_SHAPE [INPUT_SHAPE ...]]
               [--upscale_blocks UPSCALE_BLOCKS [UPSCALE_BLOCKS ...]]
               [--downscale_blocks DOWNSCALE_BLOCKS [DOWNSCALE_BLOCKS ...]]
               [--convolution_size {3}]
               --batch_size BATCH_SIZE
               [--log_every LOG_EVERY]
               [--warmup_steps WARMUP_STEPS]
               [--train_epochs TRAIN_EPOCHS]
               [--optimizer {rmsprop}]
               [--base_lr BASE_LR]
               [--momentum MOMENTUM]
               [--train_split TRAIN_SPLIT]
               [--split_seed SPLIT_SEED]
               --model_dir MODEL_DIR
               --data_dir DATA_DIR
               [--benchmark]
               [--amp]
               [--xla]
               [--augment]

Getting the data

The V-Net model was trained on the Hippocampus dataset from medical segmentation decathlon. Test images provided by the organization were used to produce the resulting masks for submission.

The objective is to produce a set of masks that segment the data as accurately as possible.

Medical segmentation decathlon (MSD) datasets are conformed by the following elements:

• dataset.json contains a high level description of the contents of the dataset
• ImagesTr contains the training images as Nifti files
• LabelsTr contains the training labels as Nifti files
• ImagesTs contains the test images as Nifti files

Dataset guidelines

The process of loading, normalizing and augmenting the data contained in the dataset can be found in the data_loader.py script.

Initially, data is loaded from a Nifti file and converted to NumPy arrays with the use of SimpleItk, with target dimensions specified through --input_shape. These NumPy arrays are fed to the model through tf.data.Dataset.from_tensor_slices(), in order to achieve high performance.

Intensities on the volumes are then normalized using the method specified in --data_normalization, whereas labels are one-hot encoded for their later use.

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

• Random horizontal flipping
• Random vertical flipping

Training process

Optimizer

The model trains for 80 epochs with the following hyperparameters:

• RMSProp optimizer with momentum = 0.0
• Base learning rate = 0.0001

Inference process

To run inference on a checkpointed model, run the script below, although it requires a pre-trained model checkpoint and tokenized input.

python examples/vnet_predict.py --data_dir ./data/Task04_Hippocampus --model_dir ./tmp --batch_size {N} [--amp]

This script should produce the prediction results over a set of masks which will be located in ./tmp/eval.