|
|
3 years ago | |
|---|---|---|
| .. | ||
| img | 3 years ago | |
| pytorch | 3 years ago | |
| Dockerfile | 3 years ago | |
| NOTICE | 7 years ago | |
| README.md | 3 years ago | |
| download_dataset.sh | 3 years ago | |
| download_weights.sh | 7 years ago | |
| hashes.md5 | 3 years ago | |
| weights.md5 | 7 years ago | |
This repository provides a script and recipe to train and infer on MaskRCNN to achieve state of the art accuracy, and is tested and maintained by NVIDIA.
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 Facebook’s implementation.This model is trained with mixed precision using Tensor Cores on Volta, Turing, and the NVIDIA Ampere GPU architectures. Therefore, researchers can get results 1.3x 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.
The 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:
These techniques/optimizations improve model performance and reduce training time by a factor of 1.3x, allowing you to perform more efficient instance segmentation with no additional effort.
Other publicly available implementations of Mask R-CNN include:
Mask R-CNN builds on top of FasterRCNN adding an additional mask head for the task of image segmentation.
The architecture consists of following:
The default configuration of this model can be found at pytorch/maskrcnn_benchmark/config/defaults.py. The default hyper-parameters are as follows:
General:
Feature extractor:
Region Proposal Network (RPN):
RoI heads:
This repository implements multi-gpu and gradient accumulation to support larger batches and mixed precision support. This implementation also includes the following optimizations.
Custom CUDA kernels for:
The source files can be found under maskrcnn_benchmark/csrc/cuda.
The following features are supported by this model.
| Feature | Mask R-CNN |
|---|---|
| Native AMP | Yes |
| Native DDP | Yes |
| Native NHWC | Yes |
AMP is an abbreviation used for automatic mixed precision training.
Native DDP; Apex DDP where DDP stands for DistributedDataParallel and is used for multi-GPU training.
NHWC is the channels last memory format for tensors.
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 the 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 requires two steps:
Porting the model to use the FP16 data type where appropriate.
Adding loss scaling to preserve small gradient values.
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.
In this repository, mixed precision training is enabled by the PyTorch native AMP library. PyTorch has an automatic mixed precision module that allows mixed precision to be enabled with minimal code changes.
Automatic mixed precision can be enabled with the following code changes:
# Create gradient scaler
scaler = torch.cuda.amp.GradScaler(init_scale=8192.0)
# Wrap the forward pass in torch.cuda.amp.autocast
with torch.cuda.amp.autocast():
loss_dict = model(images, targets)
# Gradient scaling
scaler.scale(losses).backward()
scaler.step(optimizer)
scaler.update()
AMP can be enabled by setting DTYPE to float16.
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.
MLPerf Training is an ML Commons benchmark that measures how fast systems can train models to a target quality metric. MaskRCNN is one of the MLPerf training benchmarks which is improved every year. Some of the performance optimizations used in MLPerf can be introduced to this repository easily to gain significant training speedup. Here is NVIDIA's MLPerf v1.1 submission codebase.
Listed below are some of the performance optimization tricks applied to this repository:
Increasing the local batch size and applying the above tricks gives ~2x speedup for end-to-end training time on 8 DGX A100s when compared to the old implementation.
The following sections list the requirements in order to start training the Mask R-CNN model.
This repository contains Dockerfile which extends the PyTorch NGC container and encapsulates some dependencies. Aside from these dependencies, ensure you have the following components:
For more information about how to get started with NGC containers, see the following sections from the NVIDIA GPU Cloud Documentation and the Deep Learning Documentation:
For those unable to use the [framework name] NGC container, to set up the required environment or create your own container, see the versioned NVIDIA Container Support Matrix.
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.
git clone https://github.com/NVIDIA/DeepLearningExamples.git
cd DeepLearningExamples/PyTorch/Segmentation/MaskRCNN
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
cd pytorch/
bash scripts/docker/build.sh
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
bash scripts/train.sh
The train.sh script trains a model and performs evaluation on the COCO 2014 dataset. By default, the training script:
/results directory (in the container which can be mounted to a local directory).--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.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”
bash scripts/eval.sh
Model evaluation on a checkpoint can be launched by running the pytorch/scripts/eval.sh script. The script requires:
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 = TrueMIN_BBOX_MAP = <required value>MIN_MASK_MAP = <required value>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:
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.
The following sections provide greater details of the dataset, running training and inference, and the training results.
Descriptions of the key scripts and folders are provided below.
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/
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)’
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
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.
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 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:
The default feature extractor can be changed by setting CONV_BODY parameter in yaml file to any of the following:
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:
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:
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 625SOLVER.WARMUP_FACTOR 0.01When 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.
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
Benchmarking can be performed for both training and inference. Both scripts run the Mask R-CNN model using the parameters defined in configs/e2e_mask_rcnn_R_50_FPN_1x.yaml. You can specify whether benchmarking is performed in FP16, TF32 or FP32 by specifying it as an argument to the benchmarking scripts.
Training benchmarking can performed by running the script:
scripts/train_benchmark.sh <float16/tf32/float32> <number of gpus> <NHWC True/False> <Hybrid dataloader True/False>
Inference benchmarking can be performed by running the script:
scripts/inference_benchmark.sh <float16/tf32/float32> <batch_size>
The following sections provide details on how we achieved our performance and accuracy in training and inference.
Our results were obtained by running the scripts/train.sh training script in the 21.12-py3 NGC container on NVIDIA DGX A100 (8x A100 80GB) GPUs.
| GPUs | Batch size / GPU | BBOX mAP - TF32| MASK mAP - TF32 | BBOX mAP - FP16| MASK mAP - FP16 | Time to train - TF32 | Time to train - mixed precision | Time to train speedup (TF32 to mixed precision) | --| --| -- | -- | -- | -- | -- | -- | -- | 8 | 12 | 0.3765 | 0.3408 | 0.3763 | 0.3417 | 2.15 | 1.85 | 1.16x
Our results were obtained by running the scripts/train.sh training script in the PyTorch 21.12-py3 NGC container on NVIDIA DGX-1 with 8x V100 32GB GPUs.
| GPUs | Batch size / GPU | BBOX mAP - FP32| MASK mAP - FP32 | BBOX mAP - FP16| MASK mAP - FP16 | Time to train - FP32 | Time to train - mixed precision | Time to train speedup (FP32 to mixed precision) | --| --| -- | -- | -- | -- | -- | -- | -- | 8 | 12 | 0.3768 | 0.3415 | 0.3755 | 0.3403 | 5.58 | 3.37 | 1.65x
Note: Currently V100 32GB + FP32 + NHWC + Hybrid dataloader causes a slowdown. So for all V100 32GB FP32 runs hybrid dataloader and NHWC are disabled NHWC=False HYBRID=False DTYPE=float32 bash scripts/train.sh
Here, multihead loss is simply the summation of losses on the mask head and the bounding box head.
The following tables compare mAP scores across 5 different training runs with different seeds. The runs showcase consistent convergence on all 5 seeds with very little deviation.
| Config | Seed 1 | Seed 2 | Seed 3 | Seed 4 | Seed 5 | Mean | Standard Deviation |
|---|---|---|---|---|---|---|---|
| 8 GPUs, final AP BBox | 0.3764 | 0.3766 | 0.3767 | 0.3752 | 0.3768 | 0.3763 | 0.0006 |
| 8 GPUs, final AP Segm | 0.3414 | 0.3411 | 0.341 | 0.3407 | 0.3415 | 0.3411 | 0.0003 |
Our results were obtained by running the scripts/train_benchmark.sh training script in the 21.12-py3 NGC container on NVIDIA DGX A100 (8x A100 80GB) GPUs. Performance numbers in images per second were averaged over 500 iterations.
| GPUs | Batch size / GPU | Throughput - TF32 | Throughput - mixed precision | Throughput speedup (TF32 - mixed precision) | Weak scaling - TF32 | Weak scaling - mixed precision |
|---|---|---|---|---|---|---|
| 1 | 12 | 23 | 24 | 1.04 | 1 | 1 |
| 4 | 12 | 104 | 106 | 1.02 | 4.52 | 4.42 |
| 8 | 12 | 193 | 209 | 1.08 | 8.39 | 8.71 |
Our results were obtained by running the scripts/train_benchmark.sh training script in the 21.12-py3 NGC container on NVIDIA DGX-1 with (8x V100 32GB) GPUs. Performance numbers in images per second were averaged over 500 iterations.
| GPUs | Batch size / GPU | Throughput - FP32 | Throughput - mixed precision | Throughput speedup (FP32 - mixed precision) | Weak scaling - FP32 | Weak scaling - mixed precision |
|---|---|---|---|---|---|---|
| 1 | 12 | 12 | 16 | 1.33 | 1 | 1 |
| 4 | 12 | 44 | 71 | 1.61 | 3.67 | 4.44 |
| 8 | 12 | 85 | 135 | 1.59 | 7.08 | 8.44 |
Note: Currently V100 32GB + FP32 + NHWC + Hybrid dataloader causes a slowdown. So for all V100 32GB FP32 runs hybrid dataloader and NHWC are disabled bash scripts/train_benchmark.sh fp32 <number of gpus> False False
To achieve these same results, follow the steps in the Quick Start Guide.
Our results were obtained by running the scripts/inference_benchmark.sh training script in the PyTorch 21.12-py3 NGC container on NVIDIA DGX A100 (1x A100 80GB) GPU.
FP16 Inference Latency
| Batch size | Throughput Avg | Latency Avg (ms) | Latency 90% (ms) | Latency 95% (ms) | Latency 99% (ms) |
|---|---|---|---|---|---|
| 1 | 23 | 34.91 | 33.87 | 33.95 | 34.15 |
| 2 | 26 | 59.31 | 57.80 | 57.99 | 58.27 |
| 4 | 31 | 101.46 | 99.24 | 99.51 | 99.86 |
| 8 | 31 | 197.57 | 193.82 | 194.28 | 194.77 |
TF32 Inference Latency
| Batch size | Throughput Avg | Latency Avg (ms) | Latency 90% (ms) | Latency 95% (ms) | Latency 99% (ms) |
|---|---|---|---|---|---|
| 1 | 25 | 31.66 | 31.03 | 31.13 | 31.26 |
| 2 | 28 | 56.91 | 55.88 | 56.05 | 56.02 |
| 4 | 29 | 104.11 | 102.29 | 102.53 | 102.74 |
| 8 | 30 | 201.13 | 197.43 | 197.84 | 198.19 |
Our results were obtained by running the scripts/inference_benchmark.sh training script in the PyTorch 21.12-py3 NGC container on NVIDIA DGX-1 with 1x V100 32GB GPUs.
FP16 Inference Latency
| Batch size | Throughput Avg | Latency Avg (ms) | Latency 90% (ms) | Latency 95% (ms) | Latency 99% (ms) |
|---|---|---|---|---|---|
| 1 | 19 | 44.72 | 43.62 | 43.77 | 44.03 |
| 2 | 21 | 82.80 | 81.37 | 81.67 | 82.06 |
| 4 | 22 | 155.25 | 153.15 | 153.63 | 154.10 |
| 8 | 22 | 307.60 | 304.08 | 304.82 | 305.48 |
FP32 Inference Latency
| Batch size | Throughput Avg | Latency Avg (ms) | Latency 90% (ms) | Latency 95% (ms) | Latency 99% (ms) |
|---|---|---|---|---|---|
| 1 | 16 | 52.78 | 51.87 | 52.16 | 52.43 |
| 2 | 17 | 100.81 | 99.19 | 99.67 | 100.15 |
| 4 | 17 | 202.05 | 198.84 | 199.98 | 200.92 |
| 8 | 18 | 389.99 | 384.29 | 385.77 | 387.66 |
To achieve these same results, follow the steps in the Quick Start Guide.
May 2022
October 2021
July 2021
June 2020
September 2019
July 2019
March 2019
Currently V100 32GB + FP32 + NHWC + Hybrid dataloader causes a slowdown. So for all V100 32GB FP32 runs hybrid dataloader and NHWC should be disabled.