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.
• run.py: Implements the logic for training, evaluation, and inference.
• Dockerfile: Specifies the container with the basic set of dependencies to run UNet.
• requirements.txt: Set of extra requirements for running UNet.

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

• cmd_util.py: Implements the command-line arguments parsing.
• data_loader.py: Implements the data loading and augmentation.
• losses.py: Implements the losses used during training and evaluation.
• parse_results.py: Implements the intermediate results parsing.
• setup.py: Implements helper setup functions.

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

• layers.py: Defines the different blocks that are used to assemble UNet.
• unet.py: Defines the model architecture using the blocks from the layers.py script.

Other folders included in the root directory are:

• examples/: Provides examples for training and benchmarking UNet.
• 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). Modes available:
  • train - trains model from scratch.
  • 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 --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: 1000).
• --seed: Set random seed for reproducibility (default: 0).
• --weight_decay: Weight decay coefficient (default: 0.0005).
• --log_every: Log performance every n steps (default: 100).
• --learning_rate: Model's learning rate (default: 0.0001).
• --augment: Enable data augmentation (default: False).
• --benchmark: Enable performance benchmarking (default: False). 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.
• --xla: Enable accelerated linear algebra optimization (default: False).
• --amp: Enable automatic mixed precision (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 --help

The following example output is printed when running the model:

usage: main.py [-h]
              [--exec_mode {train,train_and_predict,predict,evaluate,train_and_evaluate}]
              [--model_dir MODEL_DIR] --data_dir DATA_DIR [--log_dir LOG_DIR]
              [--batch_size BATCH_SIZE] [--learning_rate LEARNING_RATE]
              [--fold FOLD]
              [--max_steps MAX_STEPS] [--weight_decay WEIGHT_DECAY]
              [--log_every LOG_EVERY] [--warmup_steps WARMUP_STEPS]
              [--seed SEED] [--augment] [--benchmark]
              [--amp] [--xla]
 
UNet-medical
 
optional arguments:
 -h, --help            show this help message and exit
 --exec_mode {train,train_and_predict,predict,evaluate,train_and_evaluate}
                       Execution mode of running the model
 --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
 --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
                       Chosen fold for cross-validation. Use None to disable
                       cross-validation
 --max_steps MAX_STEPS
                       Maximum number of steps (batches) used for training
 --weight_decay WEIGHT_DECAY
                       Weight decay coefficient
 --log_every LOG_EVERY
                       Log performance every n steps
 --warmup_steps WARMUP_STEPS
                       Number of warmup steps
 --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 UNet model uses the EM segmentation challenge dataset. Test images provided by the organization were used to produce the resulting masks for submission. The challenge's data is made available upon registration.

Training and test data are comprised of three 512x512x30 TIF volumes (test-volume.tif, train-volume.tif and train-labels.tif). Files test-volume.tif and train-volume.tif contain grayscale 2D slices to be segmented. Additionally, training masks are provided in train-labels.tif as a 512x512x30 TIF volume, where each pixel has one of two classes:

• 0 indicating the presence of cellular membrane,
• 1 corresponding to background.

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 32-bit TIF 3D image, with values between 0 (100% membrane certainty) and 1 (100% non-membrane certainty).

Dataset guidelines

The training and test datasets are given as stacks of 30 2D-images provided as a multi-page TIF that can be read using the Pillow library and NumPy (both Python packages are installed by the Dockerfile).

Initially, data is loaded from a multi-page TIF file and converted to 512x512x30 NumPy arrays with the use of Pillow. The process of loading, normalizing and augmenting the data contained in the dataset can be found in the data_loader.py script.

These NumPy arrays are fed to the model through tf.data.Dataset.from_tensor_slices(), in order to achieve high performance.

The voxel intensities then normalized to an interval [-1, 1], whereas labels are one-hot encoded for their later use in dice or pixel-wise cross-entropy loss, becoming 512x512x30x2 tensors.

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

• Random horizontal flipping
• Random vertical flipping
• Crop to a random dimension and resize to input dimension
• Random brightness shifting

In the end, images are reshaped to 388x388 and padded to 572x572 to fit the input of the network. Masks are only reshaped to 388x388 to fit the output of the network. Moreover, pixel intensities are clipped to the [-1, 1] interval.

Multi-dataset

This implementation is tuned for the EM segmentation challenge 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 NumPy arrays by loading them at the initialization, and passing them for training in slices by tf.data.Dataset.from_tensor_slices(). If you're able to fit your dataset into the memory, then convert the data into three NumPy arrays - training images, training labels, and testing images (optional). If your dataset is large, you will have to adapt the optimizer for the lazy-loading of data. For a walk-through, check the TensorFlow tf.data API guide

The performance of the model depends on the dataset size. Generally, the model should scale better for datasets containing more data. For a smaller dataset, you might experience lower performance.

Training process

The model trains for a total 6,400 batches (6,400 / number of GPUs), with the default UNet setup:

• Adam optimizer with learning rate of 0.0001.

This default parametrization is applied when running scripts from the ./examples directory and when running main.py without explicitly overriding these parameters. By default, the training is in full precision. To enable AMP, pass the --amp flag. AMP can be enabled for every mode of execution.

The default configuration minimizes a function L = 1 - DICE + cross entropy during training.

The training can be run directly without using the predefined scripts. The name of the training script is main.py. Because of the multi-GPU support, training should always be run with the Horovod distributed launcher like this:

horovodrun -np <number/of/gpus> python main.py --data_dir /data [other parameters]

Note: When calling the main.py script manually, data augmentation is disabled. In order to enable data augmentation, use the --augment flag in your invocation.

The main result of the training are checkpoints stored by default in ./results/ on the host machine, and in the /results in the container. This location can be controlled by the --model_dir command-line argument, if a different location was mounted while starting the container. In the case when the training is run in train_and_predict mode, the inference will take place after the training is finished, and inference results will be stored to the /results directory.

If the --exec_mode train_and_evaluate parameter was used, and if --fold parameter is set to an integer value of {0, 1, 2, 3, 4}, the evaluation of the validation set takes place after the training is completed. The results of the evaluation will be printed to the console.

Inference process

Inference can be launched with the same script used for training by passing the --exec_mode predict flag:

python main.py --exec_mode predict --data_dir /data --model_dir <path/to/checkpoint> [other parameters]

The script will then:

• Load the checkpoint from the directory specified by the <path/to/checkpoint> directory (/results),
• Run inference on the test dataset,
• Save the resulting binary masks in a TIF format.