Mask R-CNN for TensorFlow2

This resource is using open-source code maintained in github (see the quick-start-guide section) and available for download from NGC

Mask R-CNN is a convolution-based neural network for the task of object instance segmentation. The paper describing the model can be found here. NVIDIA's Mask R-CNN is an optimized version of Google's TPU implementation, leveraging mixed precision arithmetic using Tensor Cores while maintaining target accuracy.

This model is trained with mixed precision using Tensor Cores on Volta, Turing, and the NVIDIA Ampere GPU architectures. Therefore, researchers can get results 2.2x faster than training without Tensor Cores, while experiencing the benefits of mixed precision training. This model is tested against each NGC monthly container release to ensure consistent accuracy and performance over time.

This repository also contains scripts to interactively launch training, benchmarking and inference routines in a Docker container.

The major differences between the official implementation of the paper and our version of Mask R-CNN are as follows:

• Mixed precision support with TensorFlow AMP
• Gradient accumulation to simulate larger batches
• Custom fused CUDA kernels for faster computations

There are other publicly NVIDIA available implementations of Mask R-CNN:

• NVIDIA PyTorch implementation
• Matterport
• Tensorpack

Model architecture

Mask R-CNN builds on top of Faster R-CNN adding a mask head for the task of image segmentation.

The architecture consists of the following:

• ResNet-50 backbone with Feature Pyramid Network (FPN)
• Region proposal network (RPN) head
• RoI Align
• Bounding and classification box head
• Mask head

Figure 1. Diagram of Mask R-CNN framework from original paper

Default configuration

The Mask R-CNN configuration and the hyper-parameters for training and testing purposes are in separate files.

The default configuration of this model can be found at mrcnn_tf2/config.py.

The default configuration is as follows:

• Feature extractor:
  • Images resized with aspect ratio maintained and smaller side length between [832,1344]
  • Ground Truth mask size 112
  • Backbone network weights are frozen after second epoch
• RPN:
  • Anchor stride set to 16
  • Anchor sizes set to (32, 64, 128, 256, 512)
  • Foreground IOU Threshold set to 0.7, Background IOU Threshold set to 0.3
  • RPN target fraction of positive proposals set to 0.5
  • Train Pre-NMS Top proposals set to 2000 per FPN layer
  • Train Post-NMS Top proposals set to 1000
  • Test Pre-NMS Top proposals set to 1000 per FPN layer
  • Test Post-NMS Top proposals set to 1000
  • RPN NMS Threshold set to 0.7
• RoI heads:
  • Foreground threshold set to 0.5
  • Batch size per image set to 512
  • A positive fraction of batch set to 0.25

The default hyper-parameters can be found at mrcnn_tf2/arguments.py. These hyperparameters can be overridden through the command-line options, in the launch scripts.

Feature support matrix

The following features are supported by this model:

| Automatic mixed precision (AMP) | Yes | | Accelerated Linear Algebra (XLA)| Yes |

Features

Automatic Mixed Precision (AMP)
This implementation of Mask-RCNN uses AMP to implement mixed precision training. It allows us to use FP16 training with FP32 master weights by modifying just a few lines of code.

XLA support (experimental)
XLA is a domain-specific compiler for linear algebra that can accelerate TensorFlow models with potentially no source code changes. The results are improvements in speed and memory usage: most internal benchmarks run ~1.1-1.5x faster after XLA is enabled.

Mixed precision training

Mixed precision is the combined use of different numerical precisions in a computational method. Mixed precision training offers significant computational speedup by performing operations in half-precision format while storing minimal information in single-precision to retain as much information as possible in critical parts of the network. Since the introduction of Tensor Cores in Volta, and following with both the Turing and Ampere architectures, significant training speedups are experienced by switching to mixed precision -- up to 3x overall speedup on the most arithmetically intense model architectures. Using mixed precision training previously required two steps:

1. Porting the model to use the FP16 data type where appropriate.
2. Adding loss scaling to preserve small gradient values.

This can now be achieved using Automatic Mixed Precision (AMP) for TensorFlow to enable the full mixed precision methodology in your existing TensorFlow model code. AMP enables mixed precision training on Volta, Turing, and NVIDIA Ampere GPU architectures automatically. The TensorFlow framework code makes all necessary model changes internally.

In TF-AMP, the computational graph is optimized to use as few casts as necessary and maximize the use of FP16, and the loss scaling is automatically applied inside of supported optimizers. AMP can be configured to work with the existing tf.contrib loss scaling manager by disabling the AMP scaling with a single environment variable to perform only the automatic mixed-precision optimization. It accomplishes this by automatically rewriting all computation graphs with the necessary operations to enable mixed precision training and automatic loss scaling.

For information about:

• How to train using mixed precision, see the Mixed Precision Training paper and Training With Mixed Precision documentation.
• Techniques used for mixed precision training, see the Mixed-Precision Training of Deep Neural Networks blog.
• How to access and enable AMP for TensorFlow, see Using TF-AMP from the TensorFlow User Guide.
• APEX tools for mixed precision training, see the NVIDIA Apex: Tools for Easy Mixed-Precision Training in PyTorch.

Enabling mixed precision

Mixed precision is enabled in TensorFlow by using the Automatic Mixed Precision (TF-AMP) extension which casts variables to half-precision upon retrieval, while storing variables in single-precision format. Furthermore, to preserve small gradient magnitudes in backpropagation, a loss scaling step must be included when applying gradients. In TensorFlow, loss scaling can be applied statically by using simple multiplication of loss by a constant value or automatically, by TF-AMP. Automatic mixed precision makes all the adjustments internally in TensorFlow, providing two benefits over manual operations. First, programmers need not modify network model code, reducing development and maintenance effort. Second, using AMP maintains forward and backward compatibility with all the APIs for defining and running TensorFlow models.

To enable mixed precision, you can simply add the values to the environmental variables inside your training script:

• Enable TF-AMP graph rewrite:

  os.environ["TF_ENABLE_AUTO_MIXED_PRECISION_GRAPH_REWRITE"] = "1"
  

• Enable Automated Mixed Precision:

  os.environ['TF_ENABLE_AUTO_MIXED_PRECISION'] = '1'
  

TF32

TensorFloat-32 (TF32) is the new math mode in NVIDIA A100 GPUs for handling the matrix math also called tensor operations. TF32 running on Tensor Cores in A100 GPUs can provide up to 10x speedups compared to single-precision floating-point math (FP32) on Volta GPUs.

TF32 Tensor Cores can speed up networks using FP32, typically with no loss of accuracy. It is more robust than FP16 for models which require high dynamic range for weights or activations.

For more information, refer to the TensorFloat-32 in the A100 GPU Accelerates AI Training, HPC up to 20x blog post.

TF32 is supported in the NVIDIA Ampere GPU architecture and is enabled by default.