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.
In this notebook, you will learn how to leverage the simplicity and convenience of TAO to:
This notebook shows an example use case for instance segmentation using the Train Adapt Optimize (TAO) Toolkit.
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
# 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())
}
}
# 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)
!cat ~/.tao_mounts.json
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:
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,
After setting up your virtual environment with the above requirements, install TAO pip package.
# SKIP this step IF you have already installed the TAO launcher.
!pip3 install nvidia-pyindex
!pip3 install nvidia-tao
# View the versions of the TAO launcher
!tao info
We will be using the COCO dataset for the tutorial. The following script will download COCO dataset automatically and convert it to TFRecords.
!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.
# verify
!ls -l $LOCAL_DATA_DIR
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.
# 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", ""))
!ngc registry model list ngc/tao/pretrained_instance_segmentation:*
!mkdir -p $LOCAL_EXPERIMENT_DIR/pretrained_resnet50/
# Pull pretrained model from NGC
!ngc registry model download-version ngc/tao/pretrained_instance_segmentation:resnet50 --dest $LOCAL_EXPERIMENT_DIR/pretrained_resnet50
print("Check that model is downloaded into dir.")
!ls -l $LOCAL_EXPERIMENT_DIR/pretrained_resnet50/pretrained_instance_segmentation_vresnet50
$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.!cat $LOCAL_SPECS_DIR/maskrcnn_train_resnet50.txt
!mkdir -p $LOCAL_EXPERIMENT_DIR/experiment_dir_unpruned
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
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
print('Model for each epoch:')
print('---------------------')
!ls -ltrh $LOCAL_EXPERIMENT_DIR/experiment_dir_unpruned/
%env NUM_STEP=25000
!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
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.
# Create an output directory to save the pruned model.
!mkdir -p $LOCAL_EXPERIMENT_DIR/experiment_dir_pruned
!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
!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.
!cat $SPECS_DIR/maskrcnn_retrain_resnet50.txt
!mkdir -p $LOCAL_EXPERIMENT_DIR/experiment_dir_retrain
!tao mask_rcnn train -e $SPECS_DIR/maskrcnn_retrain_resnet50.txt \
-d $USER_EXPERIMENT_DIR/experiment_dir_retrain\
-k $KEY \
--gpus 2
%env NUM_STEP=25000
!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
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.
# 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
# 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)
# 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)
# 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
# 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
# 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.
print('Exported model:')
print('------------')
!ls -lth $LOCAL_EXPERIMENT_DIR/export
# 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
# 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
print('Exported engine:')
print('------------')
!ls -lh $LOCAL_EXPERIMENT_DIR/export/
# 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
!ls -l $LOCAL_EXPERIMENT_DIR/maskrcnn_annotated_images