TAO Toolkit - PyTorch Backend

Overview

TAO Toolkit is a Python package hosted on the NVIDIA Python Package Index. It uses TAO containers from the NVIDIA GPU Accelerated Container Registry (NGC) for training, evaluation, export, and deployment workflows. The output of a TAO workflow is a trained model that can be deployed for inference on NVIDIA devices with DeepStream, TensorRT, or Triton.

This repository contains the PyTorch backend implementation for TAO Toolkit. It includes model code, dataclass-based experiment schemas, task entrypoints, Docker build tooling, release packaging, and tests for the TAO PyTorch model families.

Contents

• Repository Map
• Developer And Power-User Docs
• Supported Commands
• Quickstart
• Requirements
• Development Container
• Running Tasks
• Testing
• Updating the Base Docker
• Building a Release Container
• Troubleshooting
• Adding a New Network
• Contribution Guidelines
• License

Repository Map

Developer And Power-User Docs

The detailed developer and container-operation guides live in docs/index.md. Start there for agent onboarding, architecture, source-code workflows, testing/debugging, new-network integration, and prebuilt container usage.

Supported Commands

The package installs one console command per model family. Each command discovers its available subtasks from that model's scripts package. The default_specs subtask is added by the shared TAO entrypoint for every registered model.

Run this command after adding or removing a model entrypoint:

python tools/update_readme_supported_commands.py

For a specific model family, run the command with no subtask or with -h to see the supported subtasks:

depth_net -h
depth_net train -h

Quickstart

Start by preparing the checkout and development shell:

git lfs install
git lfs pull
git submodule update --init --recursive
source scripts/envsetup.sh

The envsetup.sh script sets NV_TAO_PYTORCH_TOP and defines the tao_pt function for launching the development container.

Launch an interactive container with the repository mounted at /tao-pt. Mount a host workspace if you want specs, datasets, checkpoints, and results to persist outside the container:

tao_pt --gpus all \
       --run_as_user \
       --volume /path/on/host/tao-workspace:/workspace \
       -- bash

Inside the container, generate a default experiment spec for a model family:

cd /tao-pt
depth_net default_specs results_dir=/workspace/specs/depth_net

Edit the generated /workspace/specs/depth_net/experiment.yaml with dataset, checkpoint, and output paths, then run a task:

depth_net train -e /workspace/specs/depth_net/experiment.yaml
depth_net evaluate -e /workspace/specs/depth_net/experiment.yaml
depth_net export -e /workspace/specs/depth_net/experiment.yaml

Most task commands accept Hydra-style overrides after the experiment spec:

depth_net train -e /workspace/specs/depth_net/experiment.yaml train.num_gpus=1 train.gpu_ids=[0]

Requirements

Hardware

Minimum system configuration:

• 8 GB system RAM
• 4 GB GPU RAM
• 8 CPU cores
• 1 NVIDIA GPU
• 100 GB SSD space

Recommended system configuration:

• 32 GB system RAM
• 32 GB GPU RAM
• 8 CPU cores
• 1 NVIDIA GPU
• 100 GB SSD space

Model-specific requirements can be higher, especially for large transformer, multimodal, 3D, and segmentation workloads.

Host Software

The development and release containers provide the Python, CUDA, PyTorch, and third-party packages used by TAO workflows. Prefer running tasks through tao_pt unless you are intentionally developing against a local environment.

Development Container

TAO Toolkit provides a base development Dockerfile at docker/Dockerfile. The container keeps CUDA, PyTorch, TensorRT, and model dependencies consistent across developer machines.

tao_pt usage:

usage: tao_pt [-h] [--gpus GPUS] [--volume VOLUME] [--env ENV]
              [--no-tty] [--mounts_file MOUNTS_FILE] [--shm_size SHM_SIZE]
              [--run_as_user] [--tag TAG] [--cached CACHED]
              [--ulimit ULIMIT] [--port PORT]

Example:

tao_pt --gpus all \
       --run_as_user \
       --volume /path/to/data:/data \
       --volume /path/to/results:/results \
       -- bash

Mounts File

Large datasets often live outside the repository. You can define reusable Docker mounts in ~/.tao_mounts.json or pass another file with --mounts_file.

Example:

{
  "Mounts": [
    {
      "source": "/path/to/your/experiments",
      "destination": "/workspace/tao-experiments"
    },
    {
      "source": "/path/to/config/files",
      "destination": "/workspace/tao-experiments/specs"
    }
  ]
}

Running Tasks

The common task form is:

<model_command> <subtask> -e <experiment.yaml> [hydra.overrides]

Examples:

dino train -e /workspace/specs/dino.yaml
dino evaluate -e /workspace/specs/dino.yaml evaluate.checkpoint=/results/dino/model.tlt
dino export -e /workspace/specs/dino.yaml export.checkpoint=/results/dino/model.tlt

The launcher validates that an experiment spec exists, configures visible GPUs, forwards Hydra overrides, streams logs, and records task telemetry. Some distributed workflows are launched with torchrun; most PyTorch Lightning workflows are launched directly with Python and manage devices through the shared training utilities.

Testing

Run tests from the repository root, preferably inside the development container:

pytest tests/cv_unit_test/depth_net
pytest tests/core
pytest -m "cv_unit and not tensorrt"

Useful targeted examples:

pytest tests/cv_unit_test/depth_net/test_export.py
pytest tests/multimodal_unit_test
pytest tests/ssl_unit_test

For static and CI-style checks, see the scripts in ci/, especially ci/run_static_tests.py, ci/test_changed_files.py, and ci/run_functional_tests.py.

To verify the generated README command table is up to date:

python tools/update_readme_supported_commands.py --check

Updating the Base Docker

Use this flow when updating CUDA, PyTorch, system dependencies, or common Python dependencies.

The base development Dockerfile is docker/Dockerfile. Python package requirements live in docker/requirements-pip*.txt, and apt package requirements live in docker/requirements-apt.txt.

The build script supports x86_64/AMD64 and ARM64. By default, it builds for linux/amd64.

cd $NV_TAO_PYTORCH_TOP/docker

# Build for x86_64/AMD64.
./build.sh --build --x86

# Build for ARM64.
./build.sh --build --arm

# Build and push both platforms.
./build.sh --build --multiplatform --push

For more build options:

./build.sh --help

The build script tags the newly built base Docker image with the current username. Test it with:

tao_pt --tag $USER -- bash

After validating the image, push it and update docker/manifest.json with the new digest:

./build.sh --build --push --x86
./build.sh --build --push --multiplatform

To force a rebuild without using cache:

./build.sh --build --push --force --x86

Building a Release Container

The release container is built on top of the TAO PyTorch base development image. The release flow builds a wheel for nvidia_tao_pytorch and installs it into the image defined by release/docker/Dockerfile.

git lfs install
git lfs pull
git submodule update --init --recursive
source scripts/envsetup.sh
cd $NV_TAO_PYTORCH_TOP
./release/docker/deploy.sh --build --wheel

Before cutting a release, confirm the wheel/package version in release/python/version.py and the release Docker tag assembled in release/docker/deploy.sh.

Troubleshooting

Adding a New Network

See Integrating a New Network in TAO PyTorch for the expected source layout, CLI registration steps, config schema patterns, task-script conventions, tests, and TAO Core / FTMS integration notes.

Contribution Guidelines

This repository is maintained as part of NVIDIA TAO Toolkit. Follow the existing task structure when adding or updating a model family:

• Add or update dataclass schemas under nvidia_tao_pytorch/config/<task>.
• Keep model code, dataloaders, scripts, experiment specs, and entrypoints under the corresponding task package.
• Use shared utilities from nvidia_tao_pytorch/core when possible.
• Add focused tests under the matching tests/* task directory.
• Update this README or task-specific documentation when commands or workflows change.

License

This project is licensed under the Apache-2.0 License.