DLFW changes

docker/WORKSPACE.ngc

@@ -9,24 +9,28 @@ http_archive(
         sha256 = "778197e26c5fbeb07ac2a2c5ae405b30f6cb7ad1f5510ea6fdac03bded96cc6f",
     )
 
-load("@rules_python//python:pip.bzl", "pip_install")
+load("@rules_python//python:repositories.bzl", "py_repositories")
+
+py_repositories()
 
 http_archive(
     name = "rules_pkg",
+    sha256 = "8f9ee2dc10c1ae514ee599a8b42ed99fa262b757058f65ad3c384289ff70c4b8",
     urls = [
-    	"https://mirror.bazel.build/github.com/bazelbuild/rules_pkg/releases/download/0.4.0/rules_pkg-0.4.0.tar.gz",
-	"https://github.com/bazelbuild/rules_pkg/releases/download/0.4.0/rules_pkg-0.4.0.tar.gz",
+        "https://mirror.bazel.build/github.com/bazelbuild/rules_pkg/releases/download/0.9.1/rules_pkg-0.9.1.tar.gz",
+        "https://github.com/bazelbuild/rules_pkg/releases/download/0.9.1/rules_pkg-0.9.1.tar.gz",
     ],
-    sha256 = "038f1caa773a7e35b3663865ffb003169c6a71dc995e39bf4815792f385d837d",
 )
+
 load("@rules_pkg//:deps.bzl", "rules_pkg_dependencies")
+
 rules_pkg_dependencies()
 
-git_repository(
+http_archive(
     name = "googletest",
-    remote = "https://github.com/google/googletest",
-    commit = "703bd9caab50b139428cea1aaff9974ebee5742e",
-    shallow_since = "1570114335 -0400"
+    sha256 = "755f9a39bc7205f5a0c428e920ddad092c33c8a1b46997def3f1d4a82aded6e1",
+    strip_prefix = "googletest-5ab508a01f9eb089207ee87fd547d290da39d015",
+    urls = ["https://github.com/google/googletest/archive/5ab508a01f9eb089207ee87fd547d290da39d015.zip"],
 )
 
 # External dependency for torch_tensorrt if you already have precompiled binaries.
@@ -80,17 +84,13 @@ new_local_repository(
 #########################################################################
 # Testing Dependencies (optional - comment out on aarch64)
 #########################################################################
-pip_install(
-    name = "torch_tensorrt_py_deps",
-    requirements = "//py:requirements.txt",
-)
+load("@rules_python//python:pip.bzl", "pip_parse")
 
-pip_install(
-    name = "py_test_deps",
-    requirements = "//tests/py:requirements.txt",
+pip_parse(
+    name = "devtools_deps",
+    requirements_lock = "//:requirements-dev.txt",
 )
 
-pip_install(
-    name = "pylinter_deps",
-    requirements = "//tools/linter:requirements.txt",
-)
+load("@devtools_deps//:requirements.bzl", "install_deps")
+
+install_deps()

examples/int8/training/vgg16/requirements.txt

@@ -1,4 +1,5 @@
 tensorboard>=1.14.0
+protobuf==3.20.*
 nvidia-pyindex
 --extra-index-url https://pypi.ngc.nvidia.com
 pytorch-quantization>=2.1.2

noxfile.py

@@ -191,7 +191,7 @@ def cleanup(session):
 
 def run_base_tests(session):
     print("Running basic tests")
-    session.chdir(os.path.join(TOP_DIR, "tests/py"))
+    session.chdir(os.path.join(TOP_DIR, "tests/py/ts"))
     tests = [
         "api",
         "integrations/test_to_backend_api.py",
@@ -298,7 +298,7 @@ def run_fx_tools_tests(session):
 
 def run_model_tests(session):
     print("Running model tests")
-    session.chdir(os.path.join(TOP_DIR, "tests/py"))
+    session.chdir(os.path.join(TOP_DIR, "tests/py/ts"))
     tests = [
         "models",
     ]
@@ -311,7 +311,7 @@ def run_model_tests(session):
 
 def run_accuracy_tests(session):
     print("Running accuracy tests")
-    session.chdir(os.path.join(TOP_DIR, "tests/py"))
+    session.chdir(os.path.join(TOP_DIR, "tests/py/ts"))
     tests = []
     for test in tests:
         if USE_HOST_DEPS:
@@ -340,7 +340,7 @@ def copy_model(session):
 def run_int8_accuracy_tests(session):
     print("Running accuracy tests")
     copy_model(session)
-    session.chdir(os.path.join(TOP_DIR, "tests/py"))
+    session.chdir(os.path.join(TOP_DIR, "tests/py/ts"))
     tests = [
         "ptq/test_ptq_to_backend.py",
         "ptq/test_ptq_dataloader_calibrator.py",
@@ -356,7 +356,7 @@ def run_int8_accuracy_tests(session):
 def run_trt_compatibility_tests(session):
     print("Running TensorRT compatibility tests")
     copy_model(session)
-    session.chdir(os.path.join(TOP_DIR, "tests/py"))
+    session.chdir(os.path.join(TOP_DIR, "tests/py/ts"))
     tests = [
         "integrations/test_trt_intercompatibility.py",
         # "ptq/test_ptq_trt_calibrator.py",
@@ -370,7 +370,7 @@ def run_trt_compatibility_tests(session):
 
 def run_dla_tests(session):
     print("Running DLA tests")
-    session.chdir(os.path.join(TOP_DIR, "tests/py"))
+    session.chdir(os.path.join(TOP_DIR, "tests/py/ts"))
     tests = [
         "hw/test_api_dla.py",
     ]
@@ -383,7 +383,7 @@ def run_dla_tests(session):
 
 def run_multi_gpu_tests(session):
     print("Running multi GPU tests")
-    session.chdir(os.path.join(TOP_DIR, "tests/py"))
+    session.chdir(os.path.join(TOP_DIR, "tests/py/ts"))
     tests = [
         "hw/test_multi_gpu.py",
     ]

py/torch_tensorrt/__init__.py

@@ -85,6 +85,8 @@ def _find_lib(name: str, paths: List[str]) -> str:
 from torch_tensorrt._Device import Device  # noqa: F401
 from torch_tensorrt._enums import *  # noqa: F403
 from torch_tensorrt._Input import Input  # noqa: F401
+from torch_tensorrt.logging import *
+from torch_tensorrt.ptq import *
 from torch_tensorrt._utils import *  # noqa: F403
 from torch_tensorrt._utils import sanitized_torch_version