Quick Start Guide

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.

Build the Container:

Find the resources related to this model in ngc/resources "https://ngc.nvidia.com/catalog/resources/nvidia:unet_industrial_for_tensorflow ". In this address, you can download resources manually from top right part of the page or use wget resource command to download it.

1. Make a folder using mkdir command
2. Go to that folder using cd command
3. Use the wget command to download resources as zip file inside the folder.
4. Unzip the zip file using unzip command
5. Build the container using the Dockerfile inside this directory.

mkdir unet

cd unet

wget --content-disposition https://api.ngc.nvidia.com/v2/resources/nvidia/unet_industrial_for_tensorflow/versions/20.06.3/zip -O unet_industrial_for_tensorflow_20.06.3.zip
unzip unet_industrial_for_tensorflow_20.06.3.zip

docker build . --rm -t unet_industrial:latest

Start an interactive session in the NGC container to run preprocessing/training/inference. These are the steps that we need to do to run the container in the interactive mode:

1.Make a directory for the dataset, for example ./dataset:

mkdir <path/to/dataset/directory>

2. Make a directory for results, for example ./results :

mkdir <path/to/results/directory>

3.Start the container with nvidia-docker

docker run -it --rm
-v <path/to/dataset/directory>:/data/
-v <path/to/result/directory>:/results
-p 8888:8888 unet_industrial:latest

mkdir ./dataset
mkdir ./results
docker run -it --rm \
    --shm-size=2g --ulimit memlock=-1 --ulimit stack=67108864 \
    -v <absolute path to dataset directry that we made>:/data/ \
    -v <absolute path to results directry that we made>:/results \
    -p 8888:8888 unet_industrial:latest

Download and preprocess the dataset: DAGM2007

In order to download the dataset. You can execute the command in the next cell. 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. If you face permission problem while running any of these scripts, use chmod to add the scripts and rebuild the container.

chmod +x scriptname

For example: chmod +x download_and_preprocess_dagm2007.sh

./download_and_preprocess_dagm2007.sh /data

Download the Rest of the Dataset

Go to the address that you can see in the structure that is provided by script. You make a request for having access to private data on that link and a link will be send to the email that you provided so you can download 10 groups of dataset in your local machine. Open another terminal in your local machine and use following command in the terminal to cpoy downloaded files to your running container. docker cp /home/// :/data/raw_images/private

For example, I used following command to copy the class1 folder: docker cp /home/skouchak/Class1.zip 1877b7cc7625:/data/raw_images/private

After that, unzip the folder using: unzip /folder/path/.zip

Start Training

To run training for a default configuration (as described in the 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 script requires three parameters:

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

For example, for class 1:

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

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 script 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)

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

Advanced

The following sections provide greater details of the dataset, running training and inference, and the training results. 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

General Arguments

--exec_mode=train_and_evaluate Which execution mode to run the model into.

--iter_unit=batch Will the model be run for X batches or X epochs ?

--num_iter=2500 Number of iterations to run.

--batch_size=16 Size of each minibatch per GPU.

--results_dir=/path/to/results Directory in which to write training logs, summaries and checkpoints.

--data_dir=/path/to/dataset Directory which contains the DAGM2007 dataset.

--dataset_name=DAGM2007 Name of the dataset used in this run (only DAGM2007 is supported atm).

--dataset_classID=1 ClassID to consider to train or evaluate the network (used for DAGM).

Dataset Guideline

The UNet model was trained with the Weakly Supervised Learning for Industrial Optical Inspection (DAGM 2007) dataset. The provided data is artificially generated, but similar to real world problems. It consists of multiple data sets, each consisting of 1000 images showing the background texture without defects, and of 150 images with one labeled defect each on the background texture. The images in a single data set are very similar, but each data set is generated by a different texture model and defect model. Not all deviations from the texture are necessarily defects. The algorithm will need to use the weak labels provided during the training phase to learn the properties that characterize a defect.

Training Process

Laplace Smoothing: We use this technique in the DICE loss to improve the training efficiency. This technique consists in replacing the epsilon parameter (used to avoid dividing by zero and very small: +/- 1e-7) by 1.

Adaptive Loss: The DICE Loss is not able to provide a meaningful gradient at initialisation. This leads to a model instability which often push the model to diverge. Nonetheless, once the model starts to converge, DICE loss is able to very efficiently fully train the model. Therefore, we implemented an adaptive loss which is composed of two sub-losses:

• Binary Cross-Entropy (BCE)
• DICE Loss

The model is trained with the BCE loss until the DICE Loss reach a experimentally defined threshold (0.3). Thereafter, DICE loss is used to finish training.

Weak Labelling: This dataset is referred as weakly labelled. That means that the segmentation labels are not given at the pixel level but rather in an approximate fashion.

Performance

Benchmarking

This sections shows how to run benchmarks measuring the model performance in training and inference modes.

To benchmark the training performance, you can run one of the scripts in the ./scripts/benchmarking/ directory called ./scripts/benchmarking/UNet_trainbench{AMP}{1, 4, 8}GPU.sh.

Each of the scripts requires two parameters:

• path to the dataset as the first argument
• class ID from DAGM used (between 1-10)

For example:

cd scripts/benchmarking/
./UNet_trainbench_1GPU.sh /data 1

Inference Performance Benchmark

To benchmark the inference performance, you can run one of the scripts in the ./scripts/benchmarking/ directory called ./scripts/benchmarking/UNet_evalbench{_AMP}.sh.

Each of the scripts requires:

• Path to the dataset as the first argument
• Class ID from DAGM used (between 1-10)

cd scripts/benchmarking/
./UNet_evalbench_AMP.sh /data 1

We can have access to the jupyterlb inside the container using following command inside the container bash:

jupyter notebook --ip 0.0.0.0 --port 8888 --allow-root

Running the previous command results in a link to the jupyterlab inside the container. Go to the jupyterlab and uplaod "unet-industrial-demo" notebook in the jupyterlab. By following unet-industrial-demo notebook, you can download 10 pre-trained and fine-tuned models similar to the model that you just trained. You can test these models using images in the public folders. Public datasets haven't used for training models so they can be used for testing.

You can see the process of usig a trained model and new data to test the accuracy of that trained model. If you want to use the model that you trained instead of pretrained models, you can change the checkpoints folder path to the checkpoints of model that you just trained and test your trained model using images in the public folder.