To train your model using mixed or TF32 precision with Tensor Cores or using FP32, perform the following steps using the default parameters of the Mask R-CNN model on the COCO 2017 dataset. For the specifics concerning training and inference, see the Advanced section.

1. Clone the repository.

git clone https://github.com/NVIDIA/DeepLearningExamples.git
cd DeepLearningExamples/PyTorch/Segmentation/MaskRCNN

2. Download and preprocess the dataset.

This repository provides scripts to download and extract the COCO 2017 dataset. Data will be downloaded to the current working directory on the host and extracted to a user-defined directory

To download, verify, and extract the COCO dataset, use the following scripts:

./download_dataset.sh <data/dir>

By default, the data is organized into the following structure:

<data/dir>
  annotations/
    instances_train2017.json
    instances_val2017.json
  train2017/
    *.jpg
  val2017/
    *.jpg

3. Build the Mask R-CNN PyTorch NGC container.

cd pytorch/
bash scripts/docker/build.sh

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

After you build the container image, you can start an interactive CLI session with

bash scripts/docker/interactive.sh <path/to/dataset/>

The interactive.sh script requires that the location on the dataset is specified. For example, /home/<USER>/Detectron_PyT/detectron/lib/datasets/data/coco

5. Start training.

bash scripts/train.sh

The train.sh script trains a model and performs evaluation on the COCO 2014 dataset. By default, the training script:

• Uses 8 GPUs.
• Saves a checkpoint every 2500 iterations and at the end of training. All checkpoints, evaluation results and training logs are saved to the /results directory (in the container which can be mounted to a local directory).
• Mixed precision training with Tensor Cores is invoked by either adding --amp to the command line or DTYPE \"float16\" to the end of the above command as shown in the train script. This will override the default DTYPE configuration which is tf32 for Ampere and float32 for Volta.
• Channels last memory format can be set using the NHWC flag which is set to True by default. Disabling this flag will run training using NCHW or channels first memory format.

The scripts/train.sh script runs the following Python command:

python -m torch.distributed.launch --nproc_per_node=8 tools/train_net.py --config-file "configs/e2e_mask_rcnn_R_50_FPN_1x.yaml"  

6. Start validation/evaluation.

bash scripts/eval.sh

Model evaluation on a checkpoint can be launched by running the pytorch/scripts/eval.sh script. The script requires:

• the location of the checkpoint folder to be specified and present within/mounted to the container.
• a text file named last_checkpoint which contains the path to the latest checkpoint. This mechanism is required in order to resume training from the latest checkpoint.
• The file last_checkpoint is automatically created at the end of the training process.

By default, evaluation is performed on the test dataset once training is complete. To skip evaluation at the end of training, issue the --skip-test flag.

Additionally, to perform evaluation after every epoch and terminate training on reaching a minimum required mAP score, set

• PER_EPOCH_EVAL = True
• MIN_BBOX_MAP = <required value>
• MIN_MASK_MAP = <required value>

7. Start inference/predictions.

Model predictions can be obtained on a test dataset and a model checkpoint by running the scripts/inference.sh <config/file/path> script. The script requires:

• the location of the checkpoint folder and dataset to be specified and present within/mounted to the container.
• a text file named last_checkpoint which contains the path to the checkpoint.

For example:

bash scripts/inference.sh configs/e2e_mask_rcnn_R_50_FPN_1x.yaml

Model prediction files get saved in the <OUTPUT_DIR>/inference directory and correspond to:

bbox.json - JSON file containing bounding box predictions
segm.json - JSON file containing mask predictions
predictions.pth - All prediction tensors computed by the model in torch.save() format
coco_results.pth - COCO evaluation results in torch.save() format - if --skip-eval is not used in the script above

To perform inference and skip computation of mAP scores, issue the --skip-eval flag. Performance is reported in seconds per iteration per GPU. The benchmarking scripts can be used to extract frames per second on training and inference.