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.

   git clone https://github.com/NVIDIA/DeepLearningExamples
   cd DeepLearningExamples/TensorFlow/Segmentation/UNet_Industrial
   

2. Build the UNet TensorFlow NGC container.

   # Build the docker container
   docker build . --rm -t unet_industrial:latest
   

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

   # make a directory for the dataset, for example ./dataset
   mkdir <path/to/dataset/directory>
   # make a directory for results, for example ./results
   mkdir <path/to/results/directory>
   # start the container with nvidia-docker
   nvidia-docker run -it --rm --gpus all \
       --shm-size=2g --ulimit memlock=-1 --ulimit stack=67108864 \
       -v <path/to/dataset/directory>:/data/ \
       -v <path/to/result/directory>:/results \
       unet_industrial:latest
   

4. Download and preprocess the dataset: DAGM2007

   In order to download the dataset. You can execute the following:

   ./download_and_preprocess_dagm2007.sh /data
   

   Important Information: Some files of the dataset require an account to be downloaded, the script will invite you to download them manually and put them in the correct directory.

5. Start training.

To run training for a default configuration (as described in Default configuration, for example 1/4/8 GPUs, FP32/TF-AMP), launch one of the scripts in the ./scripts directory called ./scripts/UNet{_AMP}_{1, 4, 8}GPU.sh

Each of the scripts requires three parameters:

• path to the results directory of the model as the first argument
• path to the dataset as a second argument
• class ID from DAGM used (between 1-10)

For example, for class 1:

cd scripts/
./UNet_1GPU.sh /results /data 1

6. Run evaluation

Model evaluation on a checkpoint can be launched by running one of the scripts in the ./scripts directory called ./scripts/UNet{_AMP}_EVAL.sh.

Each of the scripts requires three parameters:

• path to the results directory of the model as the first argument
• path to the dataset as a second argument
• class ID from DAGM used (between 1-10)

For example, for class 1:

cd scripts/
./UNet_EVAL.sh /results /data 1