TAO 5.5 Release - PyTorch

TAO 5.5 Release - PyTorch

Author: Arun-George-Zachariah

README.md

@@ -218,4 +218,4 @@ In order to build a new docker, please edit the `deploy.sh` file in `$NV_TAO_PYT
 TAO Toolkit PyTorch backend is not accepting contributions as part of the TAO 5.0 release, but will be open in the future.
 
 ## <a name='License'></a>License
-This project is licensed under the [Apache-2.0](./LICENSE) License.
+This project is licensed under the [Apache-2.0](./LICENSE) License.

docker/Dockerfile

@@ -1,4 +1,4 @@
-ARG PYTORCH_BASE_IMAGE=nvcr.io/nvidia/pytorch:23.12-py3
+ARG PYTORCH_BASE_IMAGE=nvcr.io/nvidia/pytorch:24.04-py3
 FROM ${PYTORCH_BASE_IMAGE}
 
 # Ensure apt-get won't prompt for selecting options
@@ -11,7 +11,6 @@ RUN apt-get upgrade && apt-get update && \
 
 # uninstall stuff from base container
 RUN pip uninstall -y sacrebleu torchtext
-RUN pip install parametrized ninja
 # Installing custom packages in /opt.
 WORKDIR /opt
 
@@ -32,25 +31,31 @@ RUN mkdir trt_oss_src && \
     -DCMAKE_CUDA_ARCHITECTURES="53;60;61;70;75;80;86;90" \
     -DTRT_LIB_DIR=/usr/lib/x86_64-linux-gnu \
     -DTRT_BIN_DIR=`pwd`/out \
-    -DCMAKE_CUDA_COMPILER=/usr/local/cuda-12.3/bin/nvcc \
-    -DCUDNN_VERSION=8.9 && \
+    -DCMAKE_CUDA_COMPILER=/usr/local/cuda-12.4/bin/nvcc \
+    -DCUDNN_VERSION=9.1.0.70 && \
   make -j16 nvinfer_plugin nvinfer_plugin_static && \
   cp libnvinfer_plugin.so.8.6.1 /usr/lib/x86_64-linux-gnu/libnvinfer_plugin.so.8.6.1 && \
   cp libnvinfer_plugin_static.a /usr/lib/x86_64-linux-gnu/libnvinfer_plugin_static.a && \
   cd ../../../ && \
   rm -rf trt_oss_src
 
-
 COPY docker/requirements-pip.txt requirements-pip.txt
-# Forcing cython==0.29.36 for pycocotools-fix with python3.10.
-RUN pip install Cython==0.29.36
-RUN pip uninstall -y jupyterlab
-RUN pip install --upgrade pip
-RUN pip install -r requirements-pip.txt \
-  && rm requirements-pip.txt
-RUN pip install -U openmim && mim install mmengine "mmcv>=2.0.0" "mmpretrain>=1.0.0rc8"
-RUN pip install "mmsegmentation>=1.0.0" mmdeploy
-RUN pip install --upgrade setuptools
+RUN pip uninstall -y jupyterlab \
+  && pip install --upgrade pip \
+  && pip install -r requirements-pip.txt \
+  && rm requirements-pip.txt \
+  && mim install "mmengine==0.10.4" "mmpretrain>=1.0.0rc8"
+
+ENV FORCE_CUDA="1"
+RUN mkdir mmcv_src && \
+  cd mmcv_src && \
+  echo "$PWD Building MMCV..." && \
+  git clone -b v2.1.0 https://github.com/open-mmlab/mmcv.git mmcv && \
+  cd mmcv && \
+  MMCV_WITH_OPS=1 pip install -v -e . 
+
+RUN pip install "mmsegmentation==1.2.2" "mmdeploy==1.3.1" "mmdet==3.3.0" "mmdet3d==1.4.0" \
+  && pip uninstall -y opencv-python && pip install opencv-python==4.8.0.74
 COPY docker/requirements-pip-pytorch.txt requirements-pip-pytorch.txt
 RUN pip install --ignore-installed --no-deps -r requirements-pip-pytorch.txt \
   && rm requirements-pip-pytorch.txt

docker/build.sh

@@ -68,4 +68,4 @@ if [ $BUILD_DOCKER = "1" ]; then
 # Exit by printing usage.
 else
     echo "Usage: ./build.sh [--build] [--push] [--force] [--default]"
-fi
+fi

docker/manifest.json

@@ -1,5 +1,6 @@
 {
     "registry": "nvcr.io",
     "repository": "nvidia/tao/tao-toolkit",
-    "tag": "5.0.0-pyt-base"
+    "digest": "sha256:ff920080025e2dd12d5740c3d722f0630f1f84630ae6f55e56f872d4914fca03",
+    "tag": "5.5.0-pyt-base"
 }

docker/requirements-pip-odise.txt

@@ -1,19 +1,19 @@
 # ODISE
-huggingface-hub
-fvcore
-ftfy
+huggingface-hub==0.20.3
+fvcore==0.1.5.post20221221
+ftfy==6.1.3
 kornia==0.6
 diffdist==0.1
 nltk>=3.6.2
-taming-transformers-rom1504
+taming-transformers-rom1504==0.0.6
 importlib-metadata==4.11.3
-flake8-comprehensions
+flake8-comprehensions==3.14.0
 git+https://github.com/facebookresearch/detectron2.git
 git+https://github.com/openai/CLIP.git@main#egg=clip
 git+https://github.com/cocodataset/panopticapi.git
 yacs>=0.1.8
 iopath==0.1.9
-jmespath
-s3transfer
-pathspec
-black
+jmespath==0.10.0
+s3transfer==0.10.0
+pathspec==0.12.1
+black==22.3.0

docker/requirements-pip-pytorch.txt

@@ -1,13 +1,13 @@
 fairscale==0.4.12
 lpips==0.1.4
-lightning-utilities==0.8.0
-pytorch-lightning==1.8.5
+lightning-utilities==0.10.1
+pytorch-lightning==2.2.0
 pytorch_metric_learning==1.7.1
-pytorch-msssim
-thop
-timm>=0.9.6.dev0
+pytorch-msssim==1.0.0
+thop==0.1.1.post2209072238
+timm==0.9.12
 torchmetrics==0.10.3
-open-clip-torch[training]==2.23.0
+open-clip-torch[training]==2.24.0
 sentencepiece==0.1.99
-ftfy
+ftfy==6.1.3
 torch-pruning==1.2.2

docker/requirements-pip.txt

@@ -1,89 +1,102 @@
-addict
-anyconfig
-astroid==2.5.2
-boto3
-botocore
-ccimport
-click
-colored
-cumm-cu114
-cutex
-easydict
-einops
-faiss-cpu
-fire
-flake8==6.0.0
-gdown
-gradio
+addict==2.4.0
+anyconfig==0.14.0
+apispec==6.6.1
+apispec_webframeworks==1.1.0
+astroid==3.0.3
+boto3==1.34.45
+botocore==1.34.45
+ccimport==0.4.2
+click==8.1.7
+colored==2.2.4
+cumm-cu114==0.4.11
+cutex==0.3.8
+Cython==0.29.36
+easydict==1.12
+einops==0.7.0
+faiss-cpu==1.7.4
+fire==0.5.0
+flake8==7.0.0
+flask==3.0.3
+flask_limiter==3.6.0
+gdown==5.1.0
+gradio==4.19.2
 hydra-core==1.2.0
-imgaug
-imageio
-isort
-lark
-lazy-import
-lazy_object_proxy
-lazy_loader
-librosa
-lmdb
-matplotlib
-mccabe
-mypy
-mypy-extensions
-natsort
-ninja
-nltk
-nvidia-eff
-nvidia-eff-tao-encryption
-numpy
-omegaconf
+imgaug==0.4.0
+imageio==2.34.0
+isort==4.3.21
+jsonlines==4.0.0
+lark==1.1.9
+lazy-import==0.2.2
+lazy-object-proxy==1.10.0
+lazy_loader==0.3
+librosa==0.10.1
+lmdb==1.4.1
+marshmallow==3.21.2
+marshmallow_enum==1.5.1
+matplotlib==3.8.2
+mccabe==0.7.0
+mypy==1.8.0
+mypy-extensions==1.0.0
+natsort==8.4.0
+ninja==1.11.1.1
+nltk==3.8.1
+nvidia-eff==0.6.5
+nvidia-eff-tao-encryption==0.1.8
+numpy==1.24.4
+omegaconf==2.3.0
 # Install onnx-graphsurgeon with extra-index-url
 --extra-index-url https://pypi.ngc.nvidia.com
-onnx-graphsurgeon
-onnx-simplifier
-onnxoptimizer
-onnxruntime
-onnxsim
+# onnx overrides 24.04 stack version (1.16.0)
+onnx==1.15.0
+onnx-graphsurgeon==0.3.27
+onnx-simplifier==0.4.35
+onnxoptimizer==0.3.13
+onnxruntime==1.17.0
+onnxsim==0.4.35
 opencv-python==4.8.0.74
-pccm
+openmim==0.3.9
+parameterized==0.9.0
+pccm==0.4.11
 pillow==9.5.0
-Polygon3
-protobuf
-pyarmor
-pyclipper
+Polygon3==3.0.9.1
+protobuf==3.20.2
+pyarmor==7.7.4
+pyclipper==1.3.0.post5
 pycocotools==2.0.6
-pycodestyle==2.10.0
-pycuda
-pyDeprecate
-pydocstyle==3.0.0
-pyflakes
-pylint==2.2.2
-pynini==2.1.5
-pyquaternion
-pyrr
-PyWavelets
-PyYAML
-rich
-scikit-image
-shapely
-soundfile
-spconv-cu114
-tabulate
-tensorboardX
-terminaltables
-tifffile
+pycodestyle==2.11.1
+pycuda==2024.1
+pyDeprecate==0.3.2
+pydocstyle==6.3.0
+pyflakes==3.2.0
+pylint==3.0.3
+pynini==2.1.6
+pyquaternion==0.9.9
+pyrr==0.10.3
+PyWavelets==1.5.0
+PyYAML==6.0.1
+rich==13.4.2
+scikit-image==0.22.0
+setuptools==69.1.0
+shapely==2.0.3
+soundfile==0.12.1
+spconv-cu114==2.3.6
+tabulate==0.9.0
+tensorboardX==2.6.2.2
+terminaltables==3.1.10
+tifffile==2024.2.12
 # Upgrading transformer due to an error with importlib version checks.
-transformers
-tokenizers
+transformers==4.37.2
+tokenizers==0.15.2
 # Same issue with tqdm.pip
-tqdm
-ujson
-unidecode
-wandb
-wget
-wrapt
-yapf
+tqdm==4.65.2
+ujson==5.9.0
+Unidecode==1.3.8
+wandb==0.16.3
+wget==3.2
+wrapt==1.16.0
+yapf==0.40.2
 # updating base py packages versions to resolve oss vulnerabilities
-starlette
-python-multipart
-fastapi
-pyarrow==14.0.1
+starlette==0.36.3
+python-multipart==0.0.9
+fastapi==0.109.2
+pyarrow==14.0.1

nvidia_tao_pytorch/api/__init__.py

@@ -0,0 +1,56 @@
+import requests
+import json
+
+def invoke_microservices(request_dict):
+    url = f"http://localhost:8000/api/v1"
+    api_endpoint = request_dict.get('api_endpoint', None)
+    neural_network_name = request_dict.get('neural_network_name', None)
+    ngc_api_key = request_dict.get('ngc_api_key', None)
+    action_name = request_dict.get('action_name', None)
+    storage = request_dict.get('storage', None)
+    specs = request_dict.get('specs', None)
+    job_id = request_dict.get('job_id', None)
+    
+    telemetry_opt_out = request_dict.get('telemetry_opt_out', "no")
+    use_ngc_staging = request_dict.get('use_ngc_staging', "True")
+    tao_api_ui_cookie = request_dict.get('tao_api_ui_cookie', "")
+    tao_api_admin_key = request_dict.get('tao_api_admin_key', "")
+    tao_api_base_url = request_dict.get('tao_api_base_url', "https://nvidia.com")
+    tao_api_status_callback_url = request_dict.get('tao_api_status_callback_url', "https://nvidia.com")
+    automl_experiment_number = request_dict.get('automl_experiment_number', "")
+    hosted_service_interaction = request_dict.get('hosted_service_interaction', "")
+
+
+    if api_endpoint == "get_networks":
+        response = requests.get(f"{url}/neural_networks")
+    elif api_endpoint == "get_actions":
+        response = requests.get(f"{url}/neural_networks/{neural_network_name}/actions")
+    elif api_endpoint == "list_ptms":
+        req_obj = {"ngc_api_key": ngc_api_key}
+        response = requests.post(f"{url}/neural_networks/{neural_network_name}/pretrained_models", req_obj)
+    elif api_endpoint == "get_schema":
+        response = requests.get(f"{url}/neural_networks/{neural_network_name}/actions/{action_name}:schema")
+    elif api_endpoint == "post_action":
+        req_obj = {"specs": specs,
+                   "cloud_metadata": storage,
+                   "ngc_api_key": ngc_api_key,
+                   "job_id": job_id,
+                   "telemetry_opt_out": telemetry_opt_out,
+                   "use_ngc_staging": use_ngc_staging,
+                   "tao_api_ui_cookie": tao_api_ui_cookie,
+                   "tao_api_admin_key": tao_api_admin_key,
+                   "tao_api_base_url": tao_api_base_url,
+                   "tao_api_status_callback_url": tao_api_status_callback_url,
+                   "automl_experiment_number": automl_experiment_number,
+                   "hosted_service_interaction": hosted_service_interaction,
+                   }
+        response = requests.post(f"{url}/neural_networks/{neural_network_name}/actions/{action_name}", data=json.dumps(req_obj))
+    elif api_endpoint == "get_jobs":
+        response = requests.get(f"{url}/neural_networks/{neural_network_name}/actions/{action_name}:ids")
+    elif api_endpoint == "get_job_status":
+        response = requests.get(f"{url}/neural_networks/{neural_network_name}/actions/{action_name}/{job_id}")
+    
+    if response and response.status_code in (200, 201):
+        return response.json()
+    else:
+        raise ValueError(f"{response.json()['error_desc']}" if response.json().get('error_desc') else f"Failed to get execute (Status Code: {response.status_code} : {response.json()})")