tests(//py): Python tests

Signed-off-by: Naren Dasan <naren@narendasan.com>
Signed-off-by: Naren Dasan <narens@nvidia.com>

Author: narendasan

WORKSPACE

@@ -3,19 +3,17 @@ workspace(name = "TRTorch")
 load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
 load("@bazel_tools//tools/build_defs/repo:git.bzl", "git_repository")
 
-
-
-
-http_archive(
+git_repository(
     name = "rules_python",
-    url = "https://github.com/bazelbuild/rules_python/releases/download/0.0.1/rules_python-0.0.1.tar.gz",
-    sha256 = "aa96a691d3a8177f3215b14b0edc9641787abaaa30363a080165d06ab65e1161",
+    remote = "https://github.com/bazelbuild/rules_python.git",
+    commit = "4fcc24fd8a850bdab2ef2e078b1de337eea751a6",
+    shallow_since = "1589292086 -0400"
 )
 
 load("@rules_python//python:repositories.bzl", "py_repositories")
 py_repositories()
-# Only needed if using the packaging rules.
-load("@rules_python//python:pip.bzl", "pip_repositories", "pip_import")
+
+load("@rules_python//python:pip.bzl", "pip_repositories", "pip3_import")
 pip_repositories()
 
 http_archive(
@@ -35,7 +33,7 @@ new_local_repository(
 )
 
 http_archive(
-    name = "libtorch_non_cxx11_abi",
+    name = "libtorch_pre_cxx11_abi",
     build_file = "@//third_party/libtorch:BUILD",
     strip_prefix = "libtorch",
     sha256 = "ea8de17c5f70015583f3a7a43c7a5cdf91a1d4bd19a6a7bc11f074ef6cd69e27",
@@ -50,6 +48,22 @@ http_archive(
     sha256 = "0efdd4e709ab11088fa75f0501c19b0e294404231442bab1d1fb953924feb6b5"
 )
 
+pip3_import(
+    name = "trtorch_py_deps",
+    requirements = "//py:requirements.txt"
+)
+
+load("@trtorch_py_deps//:requirements.bzl", "pip_install")
+pip_install()
+
+pip3_import(
+    name = "py_test_deps",
+    requirements = "//tests/py:requirements.txt"
+)
+
+load("@py_test_deps//:requirements.bzl", "pip_install")
+pip_install()
+
 # Downloaded distributions to use with --distdir
 http_archive(
     name = "cudnn",

tests/py/BUILD

@@ -0,0 +1,13 @@
+package(default_visibility = ["//visibility:public"])
+load("@py_test_deps//:requirements.bzl", "requirement")
+
+
+py_test(
+    name = "test_api",
+    srcs = [
+        "test_api.py"
+    ],
+    deps = [
+        requirement("torchvision")
+    ]
+)

tests/py/requirements.txt

@@ -0,0 +1 @@
+torchvision==0.6.0

tests/py/test_api.py

@@ -0,0 +1,81 @@
+import unittest
+import trtorch
+import torch
+import torchvision.models as models
+
+
+class ModelTestCase(unittest.TestCase):
+    def __init__(self, methodName='runTest', model=None):
+        super(ModelTestCase, self).__init__(methodName)
+        self.model = model
+        self.model.eval().to("cuda")
+
+    @staticmethod
+    def parametrize(testcase_class, model=None):
+        testloader = unittest.TestLoader()
+        testnames = testloader.getTestCaseNames(testcase_class)
+        suite = unittest.TestSuite()
+        for name in testnames:
+            suite.addTest(testcase_class(name, model=model))
+        return suite
+
+class TestCompile(ModelTestCase):
+    def setUp(self):
+        self.input = torch.randn((1, 3, 224, 224)).to("cuda")
+        self.traced_model = torch.jit.trace(self.model, [self.input])
+        self.scripted_model = torch.jit.script(self.model)
+
+    def test_compile_traced(self):
+        extra_info = {
+            "input_shapes": [self.input.shape],
+        }
+
+        trt_mod = trtorch.compile(self.traced_model, extra_info)
+        same = (trt_mod(self.input) - self.traced_model(self.input)).abs().max()
+        self.assertTrue(same < 2e-3)
+
+    #def test_compile_script(self):
+    #    pass
+
+class TestCheckMethodOpSupport(unittest.TestCase):
+    def setUp(self):
+        module = models.alexnet(pretrained=True).eval().to("cuda")
+        self.module = torch.jit.trace(module, torch.ones((1, 3, 224, 224)).to("cuda"))
+
+    def test_check_support(self):
+        self.assertTrue(trtorch.check_method_op_support(self.module, "forward"))
+
+class TestLoggingAPIs(unittest.TestCase):
+    def test_logging_prefix(self):
+        new_prefix = "TEST"
+        trtorch.logging.set_logging_prefix(new_prefix)
+        logging_prefix = trtorch.logging.get_logging_prefix()
+        self.assertEqual(new_prefix, logging_prefix)
+
+    def test_reportable_log_level(self):
+        new_level = trtorch.logging.Level.Warning
+        trtorch.logging.set_reportable_log_level(new_level)
+        level = trtorch.logging.get_reportable_log_level()
+        self.assertEqual(new_level, level)
+
+    def test_is_colored_output_on(self):
+        trtorch.logging.set_is_colored_output_on(True)
+        color = trtorch.logging.get_is_colored_output_on()
+        self.assertTrue(color)
+
+def test_suite():
+    suite = unittest.TestSuite()
+    suite.addTest(TestCompile.parametrize(TestCompile, model=models.resnet18(pretrained=True)))
+    suite.addTest(TestCompile.parametrize(TestCompile, model=models.resnet50(pretrained=True)))
+    suite.addTest(TestCompile.parametrize(TestCompile, model=models.mobilenet_v2(pretrained=True)))
+    suite.addTest(unittest.makeSuite(TestCheckMethodOpSupport))
+    suite.addTest(unittest.makeSuite(TestLoggingAPIs))
+
+    return suite
+
+suite = test_suite()
+
+runner = unittest.TextTestRunner()
+result = runner.run(suite)
+
+exit(int(not result.wasSuccessful()))