Mask R-CNN for TensorFlow2

NVIDIA Deep Learning Examples

Resource

NVIDIA Deep Learning Examples

Mask R-CNN for TensorFlow2

Mask R-CNN is a convolution based network for object instance segmentation. This implementation provides 1.3x faster training while maintaining target accuracy.

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.

Clone the repository.

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

Build the Mask R-CNN TensorFlow NGC container.

nvidia-docker build -t nvidia_mrcnn_tf2 .

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

Run the following command to launch the Docker container.
```
docker run --gpus all -it --rm --shm-size=2g --ulimit memlock=-1 --ulimit stack=67108864 nvidia_mrcnn_tf2
```
If you want to reuse the dataset and pretrained ResNet-50 weights between runs, (recommended), use -v [data directory]:/data -v [weights directory]:/weights to mount your directories inside the container:
```
docker run --gpus all -it --rm --shm-size=2g --ulimit memlock=-1 --ulimit stack=67108864 -v [data directory]:/data -v [weights directory]:/weights nvidia_mrcnn_tf2
```
The contents of /data and /weights will be downloaded in the following steps.
Download and preprocess the dataset.

This repository provides scripts to download and extract the COCO 2017 dataset.
If you already have the data, then you do not need to run the following script; instead proceed to the next step. Data will be downloaded to the [data directory] directory provided in step 3.
```
cd dataset
bash download_and_preprocess_coco.sh /data
```
Download the pre-trained ResNet-50 weights.

This repository also provides scripts to download the pre-trained weights of ResNet-50 backbone. The following script will download the pre-trained weights to /weights.
```
python scripts/download_weights.py --save_dir=/weights
```
Start training.

To run training with a default configuration (on 1/8 GPUs, AMP/FP32), run a scripts/train.py script:
```
python scripts/train.py --gpus {1,8} [--amp] 
```
The above script trains a model and evaluates the COCO 2017 dataset using the content in the /data and /weights directories. Refer to the Advanced section or run python scripts/train.py --help for more details.