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 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.
• model_fn.py: Implements the logic for training and inference.
• hooks/training_hook.py: Collects different metrics during training.
• hooks/profiling_hook.py: Collects different metrics to be used for benchmarking and testing.
• 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 --crossvalidation_idx other than None).
  • train_and_evaluate - trains model from scratch and performs validation at the end (requires --crossvalidation_idx 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).
• --crossvalidation_idx: 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]
              [--crossvalidation_idx CROSSVALIDATION_IDX]
              [--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
 --crossvalidation_idx CROSSVALIDATION_IDX
                       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