Accelerating AI Development with NVIDIA TAO Toolkit and Weights & Biases

Leveraging image classification, object detection, automatic speech recognition (ASR), and other forms of AI can fuel massive transformation within companies and business sectors. However, building AI and deep learning models from scratch is a daunting task.

A common prerequisite for building these models is having a large amount of high-quality training data and the right expertise to prepare the data, build the neural network, and continuously fine-tune the models to optimize performance. This is often a challenge for organizations beginning their machine learning (ML) journey. The high barrier to entry often becomes a blocker for enterprises of all sizes to adopt AI, despite the clear and measurable benefits AI can provide.

This post discusses how the combination of NVIDIA TAO Toolkit and the Weights & Biases MLOps platform can help kick-start an organization’s journey to leverage AI and accelerate common AI tasks. Developers can now visualize and compare multiple training runs using NVIDIA TAO (Train-Adapt-Optimize) Toolkit and Weights & Biases (W&B). In addition, we will walk through an example workflow for building an object detection model.

NVIDIA TAO Toolkit

NVIDIA TAO Toolkit is a low-code solution enabling developers and enterprises to accelerate model training and optimization processes. NVIDIA TAO Toolkit reduces the barrier to entry for anyone starting with AI by abstracting away the complexity of AI models and deep learning frameworks. With NVIDIA TAO Toolkit, you can use the power of transfer learning to fine-tune NVIDIA pretrained models with your own data and optimize the model for inference to fit the needs of your business.

NVIDIA TAO Toolkit supports a wide range of computer vision tasks such as classification, object detection, segmentation, key point estimation, OCR, and more. It provides turnkey inference optimization that reduces the model complexity and size and increases inference throughput.

Weights & Biases

Weights & Biases MLOps platform helps ML teams build better models faster. With just a few lines of code in your notebook, you can instantly debug, compare, and reproduce your models—architecture, hyperparameters, git commits, model weights, GPU usage, datasets, and predictions—all while collaborating with your teammates.

Diagram showing the five modules of the W&B platform: artifacts, tables, experiments, sweeps, and evaluation. — *Figure 1. Modules in the Weights & Biases MLOps platform*

W&B is trusted by more than 400,000 ML practitioners from some of the most innovative companies and research organizations in the world. To try it for free, sign up at Weights & Biases.

Integrating NVIDIA TAO Toolkit and Weights & Biases

NVIDIA TAO Toolkit enables you to visualize all your experimentation data within Weights and Biases. Visualize, compare, and contrast multiple training runs to determine which model candidates best suit the needs of the project and which hyperparameters have the largest impact on model performance.

The integration will also show how each model training consumes the underlying hardware to ensure resources are being fully utilized. To learn more about how to configure NVIDIA TAO Toolkit and log experiments to W&B, see the NVIDIA TAO Toolkit W&B Integration documentation.

Image classification	Object detection	Segmentation
EfficientNet B0-B5	DetectNet-v2	Unet
	FasterRCNN	MaskRCNN
	RetinaNet
	YOLOv4/v4-Tiny/v3
	SSD/DSSD
	EfficientDet

Table 1. Networks NVIDIA TAO Toolkit supports that can directly log to Weights & Biases

How to leverage W&B with NVIDIA TAO Toolkit

The following section guides you through an example tuning a pretrained object detection model within NVIDIA TAO Toolkit, leveraging Weights & Biases. First, set up the toolkit and connect your W&B account to the instance. Then log your metrics from the toolkit into W&B, and leverage the experiment tracking, data visualization, and artifact tracking capabilities.

Set up a W&B account

Create a W&B account (Note that enterprise users should check with your W&B admin.)
Copy your API key

Install and set up NVIDIA TAO Toolkit

1. Download NVIDIA TAO Toolkit resources using the call below. This package contains startup scripts to install the toolkit containers and Jupyter notebooks to run NVIDIA TAO Toolkit.

wget --content-disposition https://api.ngc.nvidia.com/v2/resources/nvidia/tao/tao-getting-started/versions/4.0.0/zip -O getting_started_v4.0.0.zip
unzip -u getting_started_v4.0.0.zip  -d ./getting_started_v4.0.0 && rm -rf getting_started_v4.0.0.zip && cd ./getting_started_v4.0.0

2. There are multiple ways to train with NVIDIA TAO Toolkit. For detailed instructions on each, see the TAO Toolkit Getting Started guide.

Launcher CLI: You can run TAO using a lightweight Python CLI app. The launcher basically acts as a frontend for the multiple NVIDIA TAO Toolkit containers built on both PyTorch and TensorFlow. The appropriate container is launched automatically based on the type of model you plan to use for your computer vision or conversational AI use cases.
Directly from container: You also have the option to run NVIDIA TAO Toolkit directly using the Docker container. To use the containers directly, you need to know which container to pull. There are multiple containers under TAO Toolkit. Depending on the model that you want to train, you will need to pull the appropriate container. This is not required when using the Launcher CLI.
TAO APIs: You can also run using REST APIs using the NVIDIA TAO Toolkit API, a Kubernetes service that enables building AI models using REST APIs. The API service can be installed on a Kubernetes cluster (local / AWS EKS) using a Helm chart along with minimal dependencies.

Configure NVIDIA TAO Toolkit to log to W&B

1. Configure the W&B API key: To communicate the data from the local compute unit and render data on the W&B server dashboard, the W&B client in the NVIDIA TAO Toolkit container must be logged in and synchronized with your profile. To include the W&B client in the container log in, set the WANDB_API_KEY environment variable in the NVIDIA TAO Toolkit containers with the API key you received when setting up your W&B account. For additional resources, see the TAO Toolkit W&B Integration documentation.

Launcher CLI: If you are using the launcher CLI, use the sample json snippet below and replace the value field under the Envs element with your W&B API key in ~/.tao_mounts.json file.

{
    "Mounts": [
        {
            "source": "/path/to/your/data",
            "destination": "/workspace/tao-experiments/data"
        },
        {
            "source": "/path/to/your/local/results",
            "destination": "/workspace/tao-experiments/results"
        },
        {
            "source": "/path/to/config/files",
            "destination": "/workspace/tao-experiments/specs"
        }
    ],
    "Envs": [
        {
            "variable": "WANDB_API_KEY",
            "value": "<api_key_value_from_wandb>"
        }
    ],
    "DockerOptions": {
        "shm_size": "16G",
        "ulimits": {
            "memlock": -1,
            "stack": 67108864
        },
        "ports": {
            "8888": 8888
        }
    }
}

Directly from Docker: If you are running directly through the container, use the -e flag of the docker command. For example, to run detectnet_v2 with W&B directly through the container, use the following script:

docker run -it --rm --gpus all \
           -v /path/in/host:/path/in/docker \
           -e WANDB_API_KEY=<api_key_value>
           nvcr.io/nvidia/tao/tao-toolkit:4.0.0-tf1.15.5 \
           detectnet_v2 train -e /path/to/experiment/spec.txt \
           -r /path/to/results/dir \
           -k $KEY --gpus 4

NVIDIA TAO APIs: For the API, add the following snippet to the values.yaml file in the TAO Toolkit Helm chart:

# Optional MLOPS setting for Weights And Biases
wandbApiKey: <wandb_api_key>

2. Run NVIDIA TAO Toolkit and log your metrics to W&B.

Instantiate the detectnet_v2.ipynb notebook present in the downloaded samples directory: notebooks/tao_launcher_starter_kit/detectnet_v2/detectnet_v2.ipynb.

Uncomment the first cell under the MLOPS integration section using the following script:

# Uncomment this code block to enable wandb integration with TAO.
   # # you can get the wandb api key by logging into https://wandb.ai
   # !pip install wandb
   #
   # os.environ[\"WANDB_API_KEY\"] = FIXME
   # import wandb
   # WANDB_LOGGED_IN = wandb.login()
   # if WANDB_LOGGED_IN:
   #     print(\"WANDB successfully logged in.

Replace the FIXME in the following line with the copied W&B key.

os.environ["WANDB_API_KEY"] = FIXME

W&B logging will be enabled by default once you pass in your API key. For more detailed information about integrating NVIDIA TAO Toolkit and W&B, see the NVIDIA TAO Toolkit W&B Integration documentation.

Figure 2 shows an example of the dashboard that automatically generates from running the notebook.

Screenshot showing the W&B experiment tracking dashboard for the computer vision example. It displays the model training metrics and a table of labeled training data, including the bounding boxes, used for the model training. — *Figure 2. Screenshot of the W&B experiment tracking dashboard for the sample computer vision project trained with NVIDIA TAO Toolkit*

Conclusion

NVIDIA TAO Toolkit helps organizations of all sizes to quickly begin leveraging AI to improve their business and operations. With the NVIDIA TAO Toolkit 4.0 release, you can use the Weights & Biases platform to track and understand how each incremental tuning job improves the performance of the AI models. For more detailed information about integrating NVIDIA TAO Toolkit and W&B, see the NVIDIA TAO Toolkit W&B Integration documentation.

Download NVIDIA TAO Toolkit and learn more. To try Weights & Biases, check out the introductory Google Colab or sign up for a W&B account.