To train your model using mixed precision with Tensor Cores or using FP32, perform the following steps using the default parameters of the UNet model on the EM segmentation challenge dataset. These steps enable you to build the UNet TensorFlow NGC container, train and evaluate your model, and generate predictions on the test data. Furthermore, you can then choose to:

• compare your evaluation accuracy with our Training accuracy results,
• compare your training performance with our Training performance benchmark,
• compare your inference performance with our Inference performance benchmark.

For the specifics concerning training and inference, see the Advanced section.

1. Clone the repository.

   Executing this command will create your local repository with all the code to run UNet.

   git clone https://github.com/NVIDIA/DeepLearningExamples
   cd DeepLearningExamples/TensorFlow2/Segmentation/UNet_Medical/
   

2. Build the UNet TensorFlow NGC container.

   This command will use the Dockerfile to create a Docker image named unet_tf2, downloading all the required components automatically.

   docker build -t unet_tf2 .
   

   The NGC container contains all the components optimized for usage on NVIDIA hardware.

3. Start an interactive session in the NGC container to run preprocessing/training/inference.

   The following command will launch the container and mount the ./data directory as a volume to the /data directory inside the container, and ./results directory to the /results directory in the container.

   mkdir data
   mkdir results
   docker run --runtime=nvidia -it --shm-size=1g --ulimit memlock=-1 --ulimit stack=67108864 --rm --ipc=host -v ${PWD}/data:/data -v ${PWD}/results:/results unet_tf2:latest /bin/bash
   

   Any datasets and experiment results (logs, checkpoints, etc.) saved to /data or /results will be accessible in the ./data or ./results directory on the host, respectively.

4. Download and preprocess the data.

   The UNet script main.py operates on data from the ISBI Challenge, the dataset originally employed in the UNet paper. The data is available to download upon registration on the website.

   Training and test data are composed of 3 multi-page TIF files, each containing 30 2D-images (around 30 Mb total). Once downloaded, the data can be used to run the training and benchmark scripts described below, by pointing main.py to its location using the --data_dir flag.

   Note: Masks are only provided for training data.

5. Start training.

   After the Docker container is launched, the training with the default hyperparameters (for example 1/8 GPUs FP32/TF-AMP) can be started with:

   bash examples/unet_TRAIN_SINGLE{_TF-AMP}.sh <number/of/gpus> <path/to/dataset> <path/to/checkpoint>
   

   For example, to run training with full precision (32-bit) on 1 GPU from the project's folder, simply use:

   bash examples/unet_TRAIN_SINGLE.sh 1 /data /results /model
   

   This script will launch a training on a single fold (fold 0) and store the model's checkpoint in the <path/to/checkpoint> directory.

   The script can be run directly by modifying flags if necessary, especially the number of GPUs, which is defined after the -np flag. Since the test volume does not have labels, 20% of the training data is used for validation in 5-fold cross-validation manner. The number of fold can be changed using --fold with an integer in range 0-4. For example, to run with 4 GPUs using fold 1 use:

   horovodrun -np 4 python main.py --data_dir /data --model_dir /results --batch_size 1 --exec_mode train --fold 1 --xla --amp
   

   Training will result in a checkpoint file being written to ./results on the host machine.

6. Start validation/evaluation.

   The trained model can be evaluated by passing the --exec_mode evaluate flag. Since evaluation is carried out on a validation dataset, the --fold parameter should be filled. For example:

   python main.py --data_dir /data --model_dir /results --batch_size 1 --exec_mode evaluate --fold 0 --xla --amp
   

   Evaluation can also be triggered jointly after training by passing the --exec_mode train_and_evaluate flag.

7. Start inference/predictions.

   The trained model can be used for inference by passing the --exec_mode predict flag:

   python main.py --data_dir /data --model_dir /results --batch_size 1 --exec_mode predict --xla --amp
   

   Now that you have your model trained and evaluated, you can choose to compare your training results with our Training accuracy results. You can also choose to benchmark the performance of your training Training performance benchmark, or Inference performance benchmark. Following the steps in these sections will ensure that you achieve the same accuracy and performance results as stated in the Results section.