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.

• mrcnn_tf2 - Contains source code of this model.
• main.py - This is the entry point that provides advanced configuration of training and evaluation processes.
• scripts/ - A folder with utility scripts that simplifies running of this model.
  • train.py - Runs training followed by evaluation.
  • evaluate.py - Runs evaluation.
  • inference.py - Runs inference.
  • benchmark_training.py - Script for running train performance benchmarks.
  • benchmark_inference.py - Script for running inference performance benchmarks.
  • download_weights.sh - Can be used to download pre-trained weights for backbone models.
• dataset/ - A folder that contains shell scripts and Python files to download the dataset.

Parameters

Below you will find a description of the most important parameters accepted by scripts. See Command-line options for list of all available options.

Utility script parameters

All the scripts in the scripts/ directory accept the following parameters:

• --batch_size N- Size of the training or evaluation batch size (depends on the script).
• --amp - When provided, enables automatic mixed precision.
• --no_xla - When provided, disables XLA (accelerated linear algebra).
• --data_dir [path] - Path to the directory that contains TFRecords of COCO 2017. Defaults to /data.
• --model_dir [path] - Output directory for information related to the model that includes model checkpoints. Defaults to /tmp/model.
• --weights_dir [path] - Directory containing pre-trained ResNet50 weights. Defaults to /weights.

Additionally, training scripts also accept some specific parameters:

• train.py
  • --gpus N - Number of GPUs to use during training.
  • --no_eval - When provided, disables evaluation after training.
• benchmark_training.py
  • --gpus N - Number of GPUs to use during training.

Note: Any additional flags not specified above will be passed to python main.py. Refer to python main.py --help for a full list of available fags.

Main script parameters

For most use cases, the scripts in scripts/ should be sufficient, but if you need more control over the model, you can also directly execute main.py.

To get the list of all parameters accepted by main.py, run python main.py --help.

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 main.py --help

The following example output is printed when running the model:

usage: main.py MODE [arguments...]

NVIDIA implementation of MastRCNN for TensorFlow 2.x

Runtime:
  MODE                                    One of supported execution modes:
                                                train - run in training mode
                                                eval - run evaluation on eval data split
                                                infer - run inference on eval data split
  --data_dir DIR                          Input directory containing the dataset (default: /data)
  --model_dir DIR                         Output directory for information related to the model (default: /results)
  --backbone_checkpoint FILE              Pretrained checkpoint for resnet (default: /weights/rn50_tf_amp_ckpt_v20.06.0/nvidia_rn50_tf_amp)
  --eval_file FILE                        Path to the validation json file (default: /data/annotations/instances_val2017.json)
  --epochs EPOCHS                         Number of training epochs (default: 12)
  --steps_per_epoch STEPS_PER_EPOCH       Number of steps (batches) per epoch. Defaults to dataset size divided by batch size. (default: None)
  --eval_samples N                        Number of evaluation samples (default: None)

Hyperparameters:
  --train_batch_size N                    Batch size (per GPU) used during training (default: 4)
  --eval_batch_size N                     Batch size used during evaluation (default: 8)
  --seed SEED                             Set a constant seed for reproducibility (default: None)
  --l2_weight_decay L2D                   Weight of l2 regularization (default: 0.0001)
  --init_learning_rate LR                 Initial learning rate (default: 0.0)
  --learning_rate_values [D [D ...]]      Learning rate decay levels that are then scaled by global batch size (default: [0.01, 0.001, 0.0001])
  --learning_rate_boundaries [N [N ...]]  Steps (in epochs) at which learning rate changes (default: [0.3, 8.0, 10.0])
  --momentum MOMENTUM                     Optimizer momentum (default: 0.9)
  --finetune_bn                           Is batchnorm finetuned training mode (default: False)
  --use_synthetic_data                    Use synthetic input data, meant for testing only (default: False)
  --xla                                   Enable XLA JIT Compiler (default: False)
  --amp                                   Enable automatic mixed precision (default: False)

Logging:
  --log_file FILE                         Output file for DLLogger logs (default: mrcnn-dlll.json)
  --log_every N                           Log performance every N steps (default: 100)
  --log_warmup_steps N                    Number of steps that will be ignored when collecting perf stats (default: 100)
  --log_graph                             Print details about TF graph (default: False)
  --log_tensorboard PATH                  When provided saves tensorboard logs to given dir (default: None)

Utility:
  -h, --help                              Show this help message and exit
  -v, --verbose                           Displays debugging logs (default: False)
  --eagerly                               Runs model in eager mode. Use for debugging only as it reduces performance. (default: False)

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 ./dataset/download_and_preprocess_coco.sh script which automatically downloads and preprocesses the training and validation sets, saving them to tfrecord files.

Dataset guidelines

The tfrecord files are fed to the model through tf.data.TFRecordDataset() to achieve high performance.

First, the images are normalized using predefined, channel-wise values (offset 0.485, 0.456, 0.406, scale 0.229, 0.224, 0.225). Then, they are augmented (random vertical flip) and resized/padded. The resizing maintains the original aspect ratio while setting the smaller side length to be between 832 and 1344.

The bounding boxes and masks are processed accordingly so that they match the processed images.

Multi-dataset

This implementation is tuned for the COCO 2017 dataset. Using other datasets is possible, but may require changes to the code (data loader) and tuning some hyperparameters (for example, learning rate, number of iterations).

In the current implementation, the data loader works with TFRecord files. If you would like to change the format of the input data, you should substitute the Dataset class which you can find in mrcnn_tf2/dataset/dataset.py.

Training process

Training is performed using the scripts/train.py script which runs main.py with the appropriate flags.

The results are displayed in the console and are saved in ./mrcnn-dll.json (can be overridden by --log_file) in a form of DLLogger logs in which you can find:

• Full configuration used during training
• Losses, learning rate and performance metrics for steps
• Final losses
• Average performance metrics

Additionally, checkpoints will be saved to /tmp/model (can be overridden by --model_dir).

Inference process

Inference is performed using the scripts/evaluate.py script which runs main.py with the appropriate flags.

The results are displayed in the console and are saved in ./mrcnn-dll.json (can be overridden by --log_file) in a form of DLLogger logs in which you can find:

• Full configuration used during the evaluation
• Evaluation metrics
• Average performance metrics