NGC | Catalog

mask_rcnn

Instance Segmentation using TAO MaskRCNN

Transfer learning is the process of transferring learned features from one application to another. It is a commonly used training technique where you use a model trained on one task and re-train to use it on a different task.

Train Adapt Optimize (TAO) Toolkit is a simple and easy-to-use Python based AI toolkit for taking purpose-built AI models and customizing them with users' own data.

Learning Objectives

In this notebook, you will learn how to leverage the simplicity and convenience of TAO to:

  • Take a pretrained resnet50 model and train a MaskRCNN model on COCO dataset
  • Evaluate the trained model
  • Run Inference with the trained model and visualize the result
  • Export the trained model to a .etlt file for deployment to DeepStream
  • Run inference on the exported. etlt model to verify deployment using TensorRT

Table of Contents

This notebook shows an example use case for instance segmentation using the Train Adapt Optimize (TAO) Toolkit.

  1. Set up env variables and map drives
  2. Installing the TAO Launcher
  3. Prepare dataset and pre-trained model
  4. Provide training specification
  5. Run TAO training
  6. Evaluate trained models
  7. Prune trained model
  8. Retrain pruned models
  9. Evaluate retrained model
  10. Visualize inferences
  11. Deploy
  12. Verify the deployed model

0. Set up env variables and map drives

When using the purpose-built pretrained models from NGC, please make sure to set the $KEY environment variable to the key as mentioned in the model overview. Failing to do so, can lead to errors when trying to load them as pretrained models.

The following notebook requires the user to set an env variable called the $LOCAL_PROJECT_DIR as the path to the users workspace. Please note that the dataset to run this notebook is expected to reside in the $LOCAL_PROJECT_DIR/data, while the TAO experiment generated collaterals will be output to $LOCAL_PROJECT_DIR/mask_rcnn. More information on how to set up the dataset and the supported steps in the TAO workflow are provided in the subsequent cells.

Note: Please make sure to remove any stray artifacts/files from the $USER_EXPERIMENT_DIR or $DATA_DOWNLOAD_DIR paths as mentioned below, that may have been generated from previous experiments. Having checkpoint files etc may interfere with creating a training graph for a new experiment.

Note: This notebook currently is by default set up to run training using 1 GPU. To use more GPU's please update the env variable $NUM_GPUS accordingly

In [1]:
# Setting up env variables for cleaner command line commands.
import os

%env KEY=nvidia_tlt
%env NUM_GPUS=1
%env USER_EXPERIMENT_DIR=/workspace/tao-experiments/mask_rcnn
%env DATA_DOWNLOAD_DIR=/workspace/tao-experiments/data

# Set this path if you don't run the notebook from the samples directory.
# %env NOTEBOOK_ROOT=~/tao-samples/mask_rcnn

# Please define this local project directory that needs to be mapped to the TAO docker session.
# The dataset expected to be present in $LOCAL_PROJECT_DIR/data, while the results for the steps
# in this notebook will be stored at $LOCAL_PROJECT_DIR/mask_rcnn
# !PLEASE MAKE SURE TO UPDATE THIS PATH!.
%env LOCAL_PROJECT_DIR=YOUR_LOCAL_EXPERIMENT_DIR

os.environ["LOCAL_DATA_DIR"] = os.path.join(
    os.getenv("LOCAL_PROJECT_DIR", os.getcwd()),
    "data"
)
os.environ["LOCAL_EXPERIMENT_DIR"] = os.path.join(
    os.getenv("LOCAL_PROJECT_DIR", os.getcwd()),
    "mask_rcnn"
)

# The sample spec files are present in the same path as the downloaded samples.
os.environ["LOCAL_SPECS_DIR"] = os.path.join(
    os.getenv("NOTEBOOK_ROOT", os.getcwd()),
    "specs"
)
%env SPECS_DIR=/workspace/tao-experiments/mask_rcnn/specs

# Showing list of specification files.
!ls -rlt $LOCAL_SPECS_DIR

The cell below maps the project directory on your local host to a workspace directory in the TAO docker instance, so that the data and the results are mapped from in and out of the docker. For more information please refer to the launcher instance in the user guide.

When running this cell on AWS, update the drive_map entry with the dictionary defined below, so that you don't have permission issues when writing data into folders created by the TAO docker.

drive_map = {
    "Mounts": [
            # Mapping the data directory
            {
                "source": os.environ["LOCAL_PROJECT_DIR"],
                "destination": "/workspace/tao-experiments"
            },
            # Mapping the specs directory.
            {
                "source": os.environ["LOCAL_SPECS_DIR"],
                "destination": os.environ["SPECS_DIR"]
            },
        ],
    "DockerOptions": {
        "user": "{}:{}".format(os.getuid(), os.getgid())
    }
}
In [2]:
# Mapping up the local directories to the TAO docker.
import json
mounts_file = os.path.expanduser("~/.tao_mounts.json")

# Define the dictionary with the mapped drives
drive_map = {
    "Mounts": [
        # Mapping the data directory
        {
            "source": os.environ["LOCAL_PROJECT_DIR"],
            "destination": "/workspace/tao-experiments"
        },
        # Mapping the specs directory.
        {
            "source": os.environ["LOCAL_SPECS_DIR"],
            "destination": os.environ["SPECS_DIR"]
        },
    ]
}

# Writing the mounts file.
with open(mounts_file, "w") as mfile:
    json.dump(drive_map, mfile, indent=4)
In [3]:
!cat ~/.tao_mounts.json

1. Installing the TAO launcher

The TAO launcher is a python package distributed as a python wheel listed in the nvidia-pyindex python index. You may install the launcher by executing the following cell.

Please note that TAO Toolkit recommends users to run the TAO launcher in a virtual env with python 3.6.9. You may follow the instruction in this page to set up a python virtual env using the virtualenv and virtualenvwrapper packages. Once you have setup virtualenvwrapper, please set the version of python to be used in the virtual env by using the VIRTUALENVWRAPPER_PYTHON variable. You may do so by running

export VIRTUALENVWRAPPER_PYTHON=/path/to/bin/python3.x

where x >= 6 and <= 8

We recommend performing this step first and then launching the notebook from the virtual environment. In addition to installing TAO python package, please make sure of the following software requirements:

  • python >=3.6.9 < 3.8.x
  • docker-ce > 19.03.5
  • docker-API 1.40
  • nvidia-container-toolkit > 1.3.0-1
  • nvidia-container-runtime > 3.4.0-1
  • nvidia-docker2 > 2.5.0-1
  • nvidia-driver > 455+

Once you have installed the pre-requisites, please log in to the docker registry nvcr.io by following the command below

docker login nvcr.io

You will be trigerred to enter a username and password. The username is $oauthtoken and the password is the API key generated from ngc.nvidia.com. Please follow the instructions in the NGC setup guide to generate your own API key.

Please note that TAO Toolkit recommends users to run the TAO launcher in a virtual env with python 3.6.9. You may follow the instruction in this page to set up a python virtual env using the virtualenv and virtualenvwrapper packages. Once you have setup virtualenvwrapper, please set the version of python to be used in the virtual env by using the VIRTUALENVWRAPPER_PYTHON variable. You may do so by running

export VIRTUALENVWRAPPER_PYTHON=/path/to/bin/python3.x where x >=6 and x <=8.

We recommend performing this step first and then launching the notebook from the virtual environment. In addition to installing TAO python package, please make sure of the following software requirements,

  • python >=3.6.9 < 3.8.x
  • docker-ce > 19.03.5
  • docker-API 1.40
  • nvidia-container-toolkit > 1.3.0-1
  • nvidia-container-runtime > 3.4.0-1
  • nvidia-docker2 > 2.5.0-1
  • nvidia-driver > 455+

After setting up your virtual environment with the above requirements, install TAO pip package.

In [4]:
# SKIP this step IF you have already installed the TAO launcher.
!pip3 install nvidia-pyindex
!pip3 install nvidia-tao
In [5]:
# View the versions of the TAO launcher
!tao info

2. Prepare dataset and pre-trained model

We will be using the COCO dataset for the tutorial. The following script will download COCO dataset automatically and convert it to TFRecords.

In [6]:
!tao mask_rcnn run bash $SPECS_DIR/download_and_preprocess_coco.sh $DATA_DOWNLOAD_DIR

Note that the dataset conversion scripts provided in specs are intended for the standard COCO dataset. If your data doesn't have caption groundtruth or test set, you can modify download_and_preprocess_coco.sh and create_coco_tf_record.py by commenting out corresponding variables.

In [7]:
# verify
!ls -l $LOCAL_DATA_DIR

Download pretrained model from NGC

We will use NGC CLI to get the pre-trained models. For more details, go to ngc.nvidia.com and click the SETUP on the navigation bar.

In [8]:
# Installing NGC CLI on the local machine.
## Download and install
%env CLI=ngccli_cat_linux.zip
!mkdir -p $LOCAL_PROJECT_DIR/ngccli

# Remove any previously existing CLI installations
!rm -rf $LOCAL_PROJECT_DIR/ngccli/*
!wget "https://ngc.nvidia.com/downloads/$CLI" -P $LOCAL_PROJECT_DIR/ngccli
!unzip -u "$LOCAL_PROJECT_DIR/ngccli/$CLI" -d $LOCAL_PROJECT_DIR/ngccli/
!rm $LOCAL_PROJECT_DIR/ngccli/*.zip 
os.environ["PATH"]="{}/ngccli:{}".format(os.getenv("LOCAL_PROJECT_DIR", ""), os.getenv("PATH", ""))
In [9]:
!ngc registry model list ngc/tao/pretrained_instance_segmentation:*
In [10]:
!mkdir -p $LOCAL_EXPERIMENT_DIR/pretrained_resnet50/
In [11]:
# Pull pretrained model from NGC
!ngc registry model download-version ngc/tao/pretrained_instance_segmentation:resnet50 --dest $LOCAL_EXPERIMENT_DIR/pretrained_resnet50
In [12]:
print("Check that model is downloaded into dir.")
!ls -l $LOCAL_EXPERIMENT_DIR/pretrained_resnet50/pretrained_instance_segmentation_vresnet50

3. Provide training specification

  • Tfrecords for the train datasets
    • In order to use the newly generated tfrecords, update the dataset_config parameter in the spec file at $SPECS_DIR/maskrcnn_train_resnet50.txt Note that the learning rate in the spec file is set for 4 GPU training. If you have N gpus, you should divide LR by 4/N.
  • Pre-trained models
  • Augmentation parameters for on the fly data augmentation
  • Other training (hyper-)parameters such as batch size, number of epochs, learning rate etc.
  • Note that the sample spec is not meant to produce SOTA accuracy on COCO. To reproduce SOTA, you might want to use TAO to train an ImageNet model first and change the total_steps to 100K or above. In one experiment, we got 37+% AP and 34% mask_AP with 8GPU training for 100K.
In [13]:
!cat $LOCAL_SPECS_DIR/maskrcnn_train_resnet50.txt

4. Train a MaskRCNN model

  • Provide the sample spec file and the output directory location for models
  • WARNING: training will take several hours or one day to complete
In [14]:
!mkdir -p $LOCAL_EXPERIMENT_DIR/experiment_dir_unpruned
In [15]:
print("For multi-GPU, change --gpus based on your machine.")
!tao mask_rcnn train -e $SPECS_DIR/maskrcnn_train_resnet50.txt \
                     -d $USER_EXPERIMENT_DIR/experiment_dir_unpruned\
                     -k $KEY \
                     --gpus 1
In [16]:
print("To resume training from a checkpoint, simply run the same training script. It will pick up from where it's left.")
!tao mask_rcnn train -e $SPECS_DIR/maskrcnn_train_resnet50.txt \
                     -d $USER_EXPERIMENT_DIR/experiment_dir_unpruned\
                     -k $KEY \
                     --gpus 1
In [17]:
print('Model for each epoch:')
print('---------------------')
!ls -ltrh $LOCAL_EXPERIMENT_DIR/experiment_dir_unpruned/

5. Evaluate trained models

In [18]:
%env NUM_STEP=25000
In [19]:
!tao mask_rcnn evaluate -e $SPECS_DIR/maskrcnn_train_resnet50.txt \
                        -m $USER_EXPERIMENT_DIR/experiment_dir_unpruned/model.step-$NUM_STEP.tlt \
                        -k $KEY

6. Prune

  • Specify pre-trained model
  • Equalization criterion (Only for resnets as they have element wise operations or MobileNets.)
  • Threshold for pruning.
  • A key to save and load the model
  • Output directory to store the model

Usually, you just need to adjust -pth (threshold) for accuracy and model size trade off. Higher pth gives you smaller model (and thus higher inference speed) but worse accuracy. The threshold value depends on the dataset and the model. 0.4 in the block below is just a start point. If the retrain accuracy is good, you can increase this value to get smaller models. Otherwise, lower this value to get better accuracy.

In [20]:
# Create an output directory to save the pruned model.
!mkdir -p $LOCAL_EXPERIMENT_DIR/experiment_dir_pruned
In [21]:
!tao mask_rcnn prune -m $USER_EXPERIMENT_DIR/experiment_dir_unpruned/model.step-$NUM_STEP.tlt \
                     -o $USER_EXPERIMENT_DIR/experiment_dir_pruned \
                     -pth 0.7 \
                     -k $KEY
In [22]:
!ls -l $LOCAL_EXPERIMENT_DIR/experiment_dir_pruned

Note that you should retrain the pruned model first, as it cannot be directly used for evaluation or inference.

7. Retrain pruned models

  • Model needs to be re-trained to bring back accuracy after pruning
  • Specify re-training specification
  • WARNING: training will take several hours or one day to complete
In [23]:
!cat $SPECS_DIR/maskrcnn_retrain_resnet50.txt
In [24]:
!mkdir -p $LOCAL_EXPERIMENT_DIR/experiment_dir_retrain
In [25]:
!tao mask_rcnn train -e $SPECS_DIR/maskrcnn_retrain_resnet50.txt \
                     -d $USER_EXPERIMENT_DIR/experiment_dir_retrain\
                     -k $KEY \
                     --gpus 2

8. Evaluate retrained model

In [26]:
%env NUM_STEP=25000
In [27]:
!tao mask_rcnn evaluate -e $SPECS_DIR/maskrcnn_retrain_resnet50.txt \
                        -m $USER_EXPERIMENT_DIR/experiment_dir_retrain/model.step-$NUM_STEP.tlt \
                        -k $KEY

9. Visualize inferences

In this section, we run the infer tool to generate inferences on the trained models and visualize the results. The infer tool produces annotated image outputs. You can choose to draw bounding boxes only or draw both bboxes and masks.

In [28]:
# Running inference for detection on n images
!tao mask_rcnn inference -i $DATA_DOWNLOAD_DIR/raw-data/test2017 \
                         -o $USER_EXPERIMENT_DIR/maskrcnn_annotated_images \
                         -e $SPECS_DIR/maskrcnn_train_resnet50.txt \
                         -m $USER_EXPERIMENT_DIR/experiment_dir_unpruned/model.step-$NUM_STEP.tlt \
                         -l $SPECS_DIR/coco_labels.txt \
                         -t 0.5 \
                         -k $KEY \
                         --include_mask
In [29]:
# Simple grid visualizer
import matplotlib.pyplot as plt
import os
from math import ceil
valid_image_ext = ['.jpg']

def visualize_images(image_dir, num_cols=4, num_images=10):
    output_path = os.path.join(os.environ['LOCAL_EXPERIMENT_DIR'], image_dir)
    num_rows = int(ceil(float(num_images) / float(num_cols)))
    f, axarr = plt.subplots(num_rows, num_cols, figsize=[80,30])
    f.tight_layout()
    a = [os.path.join(output_path, image) for image in os.listdir(output_path) 
         if os.path.splitext(image)[1].lower() in valid_image_ext]
    for idx, img_path in enumerate(a[:num_images]):
        col_id = idx % num_cols
        row_id = idx // num_cols
        img = plt.imread(img_path)
        axarr[row_id, col_id].imshow(img) 
In [30]:
# Visualizing the sample images.
OUTPUT_PATH = 'maskrcnn_annotated_images' # relative path from $USER_EXPERIMENT_DIR.
COLS = 2 # number of columns in the visualizer grid.
IMAGES = 4 # number of images to visualize.

visualize_images(OUTPUT_PATH, num_cols=COLS, num_images=IMAGES)

10. Deploy!

In [31]:
# Export in FP32 mode. 
!mkdir -p $LOCAL_EXPERIMENT_DIR/export 
!tao mask_rcnn export -m $USER_EXPERIMENT_DIR/experiment_dir_retrain/model.step-$NUM_STEP.tlt \
                      -k $KEY \
                      -e $SPECS_DIR/maskrcnn_retrain_resnet50.txt \
                      --batch_size 1 \
                      --data_type fp32 \
                      --engine_file $USER_EXPERIMENT_DIR/export/model.step-$NUM_STEP.engine
In [32]:
# Export in INT8 mode. 
!mkdir -p $LOCAL_EXPERIMENT_DIR/export
# Uncomment to remove existing etlt file
# !rm $LOCAL_EXPERIMENT_DIR/experiment_dir_retrain/model.step-25000.etlt
!tao mask_rcnn export -m $USER_EXPERIMENT_DIR/experiment_dir_retrain/model.step-$NUM_STEP.tlt \
                      -k $KEY \
                      -e $SPECS_DIR/maskrcnn_retrain_resnet50.txt \
                      --batch_size 1 \
                      --data_type int8 \
                      --cal_image_dir $DATA_DOWNLOAD_DIR/raw-data/val2017 \
                      --batches 10 \
                      --cal_cache_file $USER_EXPERIMENT_DIR/export/maskrcnn.cal \
                      --cal_data_file $USER_EXPERIMENT_DIR/export/maskrcnn.tensorfile
In [33]:
# Check if etlt model is correctly saved.
!ls -l $LOCAL_EXPERIMENT_DIR/experiment_dir_retrain/

Verify engine generation using the tao-converter utility included with the docker.

The tao-converter produces optimized tensorrt engines for the platform that it resides on. Therefore, to get maximum performance, please instantiate this docker and execute the tao-converter command, with the exported .etlt file and calibration cache (for int8 mode) on your target device. The tao-converter utility included in this docker only works for x86 devices, with discrete NVIDIA GPU's.

For the jetson devices, please download the tao-converter for jetson from the dev zone link here.

If you choose to integrate your model into deepstream directly, you may do so by simply copying the exported .etlt file along with the calibration cache to the target device and updating the spec file that configures the gst-nvinfer element to point to this newly exported model. Please refer to deepstream dev guide for more details.

In [34]:
print('Exported model:')
print('------------')
!ls -lth $LOCAL_EXPERIMENT_DIR/export
In [35]:
# Convert to TensorRT engine(FP16).
!tao converter -k $KEY  \
                   -d 3,832,1344 \
                   -o generate_detections,mask_fcn_logits/BiasAdd \
                   -e $USER_EXPERIMENT_DIR/export/trt.fp16.engine \
                   -t fp16 \
                   -i nchw \
                   -m 1 \
                   $USER_EXPERIMENT_DIR/experiment_dir_retrain/model.step-$NUM_STEP.etlt
In [36]:
# Convert to TensorRT engine(INT8).
!tao converter -k $KEY  \
                   -d 3,832,1344 \
                   -o generate_detections,mask_fcn_logits/BiasAdd \
                   -c $USER_EXPERIMENT_DIR/export/maskrcnn.cal \
                   -e $USER_EXPERIMENT_DIR/export/trt.int8.engine \
                   -b 8 \
                   -m 1 \
                   -t int8 \
                   -i nchw \
                   $USER_EXPERIMENT_DIR/experiment_dir_retrain/model.step-$NUM_STEP.etlt
In [37]:
print('Exported engine:')
print('------------')
!ls -lh $LOCAL_EXPERIMENT_DIR/export/

11. Verify the deployed model

Verify the converted engine by visualizing TensorRT inferences.

In [38]:
# Running inference for detection on a dir of images
!tao mask_rcnn inference -i $DATA_DOWNLOAD_DIR/raw-data/test2017 \
                         -o $USER_EXPERIMENT_DIR/maskrcnn_annotated_images \
                         -e $SPECS_DIR/maskrcnn_retrain_resnet50.txt \
                         -m $USER_EXPERIMENT_DIR/export/model.step-$NUM_STEP.engine \
                         -l $USER_EXPERIMENT_DIR/maskrcnn_annotated_labels \
                         -c $SPECS_DIR/coco_labels.txt \
                         -t 0.5 \
                         --include_mask
In [39]:
!ls -l $LOCAL_EXPERIMENT_DIR/maskrcnn_annotated_images