The following sections provide greater details of the dataset, running training and inference, and the training results.

Scripts and sample code

Descriptions of the key scripts and folders are provided below.

• maskrcnn_benchmark - Contains scripts to build individual components of the model such as backbone, FPN, RPN, mask and bbox heads etc.,

• download_dataset.sh - Launches download and processing of required datasets.

• scripts/ - Contains shell scripts to launch data download, train the model and perform inferences.

• train.sh - Launches model training

• eval.sh - Performs inference and computes mAP of predictions.

• inference.sh - Performs inference on given data.

• train_benchmark.sh - To benchmark training performance.

• inference_benchmark.sh - To benchmark inference performance.

• docker/ - Scripts to build the docker image and to start an interactive session.

• tools/

  • train_net.py - End to end to script to load data, build and train the model.
  • test_net.py - End to end script to load data, checkpoint and perform inference and compute mAP score.

Parameters

train_net.py script parameters

You can modify the training behaviour through the various flags in both the train_net.py script and through overriding specific parameters in the YAML config files. Flags in the train_net.py script are as follows:

--config_file - path to config file containing model params

--skip-test - skips model testing after training

--opts - allows for you to override specific params in config file

For example:

python -m torch.distributed.launch --nproc_per_node=2 tools/train_net.py \
    --config-file configs/e2e_faster_rcnn_R_50_FPN_1x.yaml \
    DTYPE "float16" \
    NHWC True \
    OUTPUT_DIR RESULTS \
    SOLVER.BASE_LR 0.002 \
    SOLVER.STEPS '(360000, 480000)'

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 tools/train_net.py --help

Getting the data

The Mask R-CNN model was trained on the COCO 2017 dataset. This dataset comes with a training and validation set.

This repository contains the ./download_dataset.sh,./verify_dataset.sh, and ./extract_dataset.sh scripts which automatically download and preprocess the training and validation sets.

Dataset guidelines

In order to run on your own dataset, ensure your dataset is present/mounted to the Docker container with the following hierarchy:

my_dataset/
  images_train/
  images_val/
  instances_train.json
  instances_val.json

and add it to DATASETS dictionary in maskrcnn_benchmark/config/paths_catalog.py

DATASETS = {
        "my_dataset_train": {
            "img_dir": "data/images_train",
            "ann_file": "data/instances_train.json"
        },
        "my_dataset_val": {
            "img_dir": "data/images_val",
            "ann_file": "data/instances_val.json"
        },
      }

python -m torch.distributed.launch --nproc_per_node=<NUM_GPUS> tools/train_net.py \
        --config-file <CONFIG? \
        DATASETS.TRAIN "(\"my_dataset_train\")"\
        DATASETS.TEST "(\"my_dataset_val\")"\
        DTYPE "float16" \
        OUTPUT_DIR <RESULTS> \
        | tee <LOGFILE>

Training Process

Training is performed using the tools/train_net.py script along with parameters defined in the config file. The default config files can be found in the pytorch/configs/ directory.

The e2e_mask_rcnn_R_50_FPN_1x.yaml file was used to gather accuracy and performance metrics. This configuration sets the following parameters:

• Backbone weights to ResNet-50
• Feature extractor set to ResNet-50 with Feature Pyramid Networks (FPN)
• RPN uses FPN
• RoI Heads use FPN
• Dataset - COCO 2017
• Base Learning Rate - 0.12
• Global train batch size - 96
• Global test batch size - 8
• RPN batch size - 256
• ROI batch size - 512
• Solver steps - (12000, 16000)
• Max iterations - 16667
• Warmup iterations - 800
• Warmup factor = 0.0001
  • Initial learning rate = Base Learning Rate x Warmup factor

The default feature extractor can be changed by setting CONV_BODY parameter in yaml file to any of the following:

• R-50-C4
• R-50-C5
• R-101-C4
• R-101-C5
• R-101-FPN

The default backbone can be changed to a flavor of Resnet-50 or ResNet-101 by setting WEIGHT parameter in yaml file to any of the following:

• "catalog://ImageNetPretrained/MSRA/R-50-GN"
• "catalog://ImageNetPretrained/MSRA/R-101"
• "catalog://ImageNetPretrained/MSRA/R-101-GN"

This script outputs results to the current working directory by default. However, this can be changed by adding OUTPUT_DIR <DIR_NAME> to the end of the default command. Logs produced during training are also stored in the OUTPUT_DIR specified. The training log will contain information about:

• Loss, time per iteration, learning rate and memory metrics
• performance values such as time per step
• test accuracy and test performance values after evaluation

The training logs are located in the <OUTPUT_DIR>/log directory. The summary after each training epoch is printed in the following format:

INFO:maskrcnn_benchmark.trainer:eta: 4:42:15  iter: 20  loss: 1.8236 (2.7274)  loss_box_reg: 0.0249 (0.0620)  loss_classifier: 0.6086 (1.2918)  loss_mask: 0.6996 (0.8026)  loss_objectness: 0.5373 (0.4787)  loss_rpn_box_reg: 0.0870 (0.0924)  time: 0.2002 (0.3765)  data: 0.0099 (0.1242)  lr: 0.014347  max mem: 3508

The mean and median training losses are reported every 20 steps.

Multi-gpu and multi-node training is enabled with the PyTorch distributed launch module. The following example runs training on 8 GPUs:

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

We have tested batch sizes upto 12 on a 32GB V100 and 80GB A100 with mixed precision. The repository also implements gradient accumulation functionality to simulate bigger batches as follows:

python -m torch.distributed.launch --nproc_per_node=8 tools/train_net.py --config-file \"configs/e2e_mask_rcnn_R_50_FPN_1x.yaml\" SOLVER.ACCUMULATE_GRAD True SOLVER.ACCUMULATE_STEPS 4

By default, training is performed using FP32 on Volta and TF32 on Ampere, however training time can be reduced further using tensor cores and mixed precision. This can be done by either adding --amp to the command line or DTYPE \"float16\" to override the respective parameter in the config file.

Note: When training a global batch size >= 32, it is recommended to add required warmup by additionally setting the following parameters:

• SOLVER.WARMUP_ITERS 625
• SOLVER.WARMUP_FACTOR 0.01

When experimenting with different global batch sizes for training and inference, make sure SOLVER.IMS_PER_BATCH and TEST.IMS_PER_BATCH are divisible by the number of GPUs.

Other training options

A sample single GPU config is provided under configs/e2e_mask_rcnn_R_50_FPN_1x_1GPU.yaml

To train with smaller global batch sizes (32 or 64) use configs/e2e_mask_rcnn_R_50_FPN_1x_bs32.yaml and configs/e2e_mask_rcnn_R_50_FPN_1x_bs64.yaml respectively.

For multi-gpu runs, -m torch.distributed.launch --nproc_per_node num_gpus is added prior to tools/train_net.py. For example, for an 8 GPU run:

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

Training is terminated when either the required accuracies specified on the command line are reached or if the number of training iterations specified is reached.

To terminate training on reaching target accuracy on 8 GPUs, run:

python -m torch.distributed.launch --nproc_per_node=8 tools/train_net.py --config-file "configs/e2e_mask_rcnn_R_50_FPN_1x.yaml" PER_EPOCH_EVAL True MIN_BBOX_MAP 0.377 MIN_MASK_MAP 0.342

Note: The score is always the Average Precision(AP) at

• IoU = 0.50:0.95
• Area = all - include small, medium and large
• maxDets = 100