chore: Linter fixes

Signed-off-by: Dheeraj Peri <[email protected]>

Author: peri044

.github/workflows/docgen.yml

@@ -31,7 +31,7 @@ jobs:
       - name: Set up Python 3.9.4
         uses: actions/setup-python@v2
         with:
-          python-version: 3.9.4 
+          python-version: 3.9.4
       - uses: actions/checkout@v2
         with:
           ref: ${{github.head_ref}}

.github/workflows/linter.yml

@@ -39,7 +39,7 @@ jobs:
           pip3 install -r $GITHUB_WORKSPACE/.github/scripts/requirements.txt
           pip3 install -r $GITHUB_WORKSPACE/requirements-dev.txt
       - name: Lint C++
-        run: | 
+        run: |
           cd $GITHUB_WORKSPACE
           python3 $GITHUB_WORKSPACE/.github/scripts/run_cpp_linter.py
         env:

noxfile.py

@@ -31,14 +31,15 @@
     print("Using dependencies from host python")
 
 # Set epochs to train VGG model for accuracy tests
-EPOCHS=25
+EPOCHS = 25
 
 SUPPORTED_PYTHON_VERSIONS = ["3.7", "3.8", "3.9", "3.10"]
 
 nox.options.sessions = [
     "l0_api_tests-" + "{}.{}".format(sys.version_info.major, sys.version_info.minor)
 ]
 
+
 def install_deps(session):
     print("Installing deps")
     session.install("-r", os.path.join(TOP_DIR, "py", "requirements.txt"))
@@ -110,7 +111,9 @@ def train_model(session):
             str(EPOCHS),
         )
 
-        session.run_always("python", "export_ckpt.py", "vgg16_ckpts/ckpt_epoch" + str(EPOCHS) + ".pth")
+        session.run_always(
+            "python", "export_ckpt.py", "vgg16_ckpts/ckpt_epoch" + str(EPOCHS) + ".pth"
+        )
 
 
 def finetune_model(session):
@@ -135,15 +138,15 @@ def finetune_model(session):
             "--start-from",
             str(EPOCHS),
             "--epochs",
-            str(EPOCHS+1),
+            str(EPOCHS + 1),
             env={"PYTHONPATH": PYT_PATH},
         )
 
         # Export model
         session.run_always(
             "python",
             "export_qat.py",
-            "vgg16_ckpts/ckpt_epoch" + str(EPOCHS+1) + ".pth",
+            "vgg16_ckpts/ckpt_epoch" + str(EPOCHS + 1) + ".pth",
             env={"PYTHONPATH": PYT_PATH},
         )
     else:
@@ -161,11 +164,15 @@ def finetune_model(session):
             "--start-from",
             str(EPOCHS),
             "--epochs",
-            str(EPOCHS+1),
+            str(EPOCHS + 1),
         )
 
         # Export model
-        session.run_always("python", "export_qat.py", "vgg16_ckpts/ckpt_epoch" + str(EPOCHS+1) + ".pth")
+        session.run_always(
+            "python",
+            "export_qat.py",
+            "vgg16_ckpts/ckpt_epoch" + str(EPOCHS + 1) + ".pth",
+        )
 
 
 def cleanup(session):
@@ -195,6 +202,7 @@ def run_base_tests(session):
         else:
             session.run_always("pytest", test)
 
+
 def run_model_tests(session):
     print("Running model tests")
     session.chdir(os.path.join(TOP_DIR, "tests/py"))
@@ -258,7 +266,7 @@ def run_trt_compatibility_tests(session):
     session.chdir(os.path.join(TOP_DIR, "tests/py"))
     tests = [
         "integrations/test_trt_intercompatibility.py",
-        #"ptq/test_ptq_trt_calibrator.py",
+        # "ptq/test_ptq_trt_calibrator.py",
     ]
     for test in tests:
         if USE_HOST_DEPS:
@@ -310,6 +318,7 @@ def run_l0_dla_tests(session):
     run_base_tests(session)
     cleanup(session)
 
+
 def run_l1_model_tests(session):
     if not USE_HOST_DEPS:
         install_deps(session)
@@ -318,6 +327,7 @@ def run_l1_model_tests(session):
     run_model_tests(session)
     cleanup(session)
 
+
 def run_l1_int8_accuracy_tests(session):
     if not USE_HOST_DEPS:
         install_deps(session)
@@ -352,16 +362,19 @@ def l0_api_tests(session):
     """When a developer needs to check correctness for a PR or something"""
     run_l0_api_tests(session)
 
+
 @nox.session(python=SUPPORTED_PYTHON_VERSIONS, reuse_venv=True)
 def l0_dla_tests(session):
     """When a developer needs to check basic api functionality using host dependencies"""
     run_l0_dla_tests(session)
 
+
 @nox.session(python=SUPPORTED_PYTHON_VERSIONS, reuse_venv=True)
 def l1_model_tests(session):
     """When a developer needs to check correctness for a PR or something"""
     run_l1_model_tests(session)
 
+
 @nox.session(python=SUPPORTED_PYTHON_VERSIONS, reuse_venv=True)
 def l1_int8_accuracy_tests(session):
     """Checking accuracy performance on various usecases"""

tests/cpp/test_collections.cpp

@@ -192,7 +192,7 @@ TEST(CppAPITests, TestCollectionTupleInputOutput) {
   auto trt_out = trt_mod.forward(complex_inputs);
 
   ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(
-        out.toTuple()->elements()[0].toTensor(), trt_out.toTuple()->elements()[0].toTensor(), 1e-5));
+      out.toTuple()->elements()[0].toTensor(), trt_out.toTuple()->elements()[0].toTensor(), 1e-5));
   ASSERT_TRUE(torch_tensorrt::tests::util::almostEqual(
       out.toTuple()->elements()[1].toTensor(), trt_out.toTuple()->elements()[1].toTensor(), 1e-5));
 }

tests/cpp/test_modules_as_engines.cpp

@@ -14,8 +14,8 @@ TEST_P(CppAPITests, ModuleAsEngineIsClose) {
   jit_results.push_back(jit_results_ivalues.toTensor());
   auto trt_results = torch_tensorrt::tests::util::RunModuleForwardAsEngine(mod, inputs);
 
-  ASSERT_TRUE(
-      torch_tensorrt::tests::util::cosineSimEqual(jit_results[0], trt_results[0].reshape_as(jit_results[0]), threshold));
+  ASSERT_TRUE(torch_tensorrt::tests::util::cosineSimEqual(
+      jit_results[0], trt_results[0].reshape_as(jit_results[0]), threshold));
 }
 
 #ifndef DISABLE_TEST_IN_CI

tests/py/api/test_collections.py

@@ -41,8 +41,13 @@ def test_compile(self):
         }
 
         trt_mod = torchtrt.ts.compile(self.model, **compile_spec)
-        cos_sim = cosine_similarity(self.model(self.input, self.input), trt_mod(self.input, self.input))
-        self.assertTrue(cos_sim > COSINE_THRESHOLD, msg=f"standard_tensor_input_scripted TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}")
+        cos_sim = cosine_similarity(
+            self.model(self.input, self.input), trt_mod(self.input, self.input)
+        )
+        self.assertTrue(
+            cos_sim > COSINE_THRESHOLD,
+            msg=f"standard_tensor_input_scripted TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}",
+        )
 
 
 class TestTupleInput(unittest.TestCase):
@@ -65,8 +70,13 @@ def test_compile(self):
         }
 
         trt_mod = torchtrt.ts.compile(self.model, **compile_spec)
-        cos_sim = cosine_similarity(self.model((self.input, self.input)), trt_mod((self.input, self.input)))
-        self.assertTrue(cos_sim > COSINE_THRESHOLD, msg=f"tuple_input_scripted TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}")
+        cos_sim = cosine_similarity(
+            self.model((self.input, self.input)), trt_mod((self.input, self.input))
+        )
+        self.assertTrue(
+            cos_sim > COSINE_THRESHOLD,
+            msg=f"tuple_input_scripted TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}",
+        )
 
 
 class TestListInput(unittest.TestCase):
@@ -87,8 +97,13 @@ def test_compile(self):
         }
 
         trt_mod = torchtrt.ts.compile(self.model, **compile_spec)
-        cos_sim = cosine_similarity(self.model([self.input, self.input]), trt_mod([self.input, self.input]))
-        self.assertTrue(cos_sim > COSINE_THRESHOLD, msg=f"list_input_scripted TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}")
+        cos_sim = cosine_similarity(
+            self.model([self.input, self.input]), trt_mod([self.input, self.input])
+        )
+        self.assertTrue(
+            cos_sim > COSINE_THRESHOLD,
+            msg=f"list_input_scripted TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}",
+        )
 
 
 class TestTupleInputOutput(unittest.TestCase):
@@ -115,7 +130,10 @@ def test_compile(self):
         pyt_out = self.model((self.input, self.input))
         for (t, p) in zip(trt_out, pyt_out):
             cos_sim = cosine_similarity(t, p)
-            self.assertTrue(cos_sim > COSINE_THRESHOLD, msg=f"tuple_input_output_scripted TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}")
+            self.assertTrue(
+                cos_sim > COSINE_THRESHOLD,
+                msg=f"tuple_input_output_scripted TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}",
+            )
 
 
 class TestListInputOutput(unittest.TestCase):
@@ -143,7 +161,10 @@ def test_compile(self):
 
         for (t, p) in zip(trt_out, pyt_out):
             cos_sim = cosine_similarity(t, p)
-            self.assertTrue(cos_sim > COSINE_THRESHOLD, msg=f"list_input_output_scripted TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}")
+            self.assertTrue(
+                cos_sim > COSINE_THRESHOLD,
+                msg=f"list_input_output_scripted TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}",
+            )
 
 
 class TestListInputTupleOutput(unittest.TestCase):
@@ -170,7 +191,10 @@ def test_compile(self):
         pyt_out = self.model((self.input, self.input))
         for (t, p) in zip(trt_out, pyt_out):
             cos_sim = cosine_similarity(t, p)
-            self.assertTrue(cos_sim > COSINE_THRESHOLD, msg=f"list_input_tuple_output_scripted TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}")
+            self.assertTrue(
+                cos_sim > COSINE_THRESHOLD,
+                msg=f"list_input_tuple_output_scripted TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}",
+            )
 
 
 if __name__ == "__main__":

tests/py/api/test_e2e_behavior.py

@@ -5,6 +5,7 @@
 import copy
 from typing import Dict
 
+
 class TestInputTypeDefaultsFP32Model(unittest.TestCase):
     def test_input_use_default_fp32(self):
         self.model = models.resnet18(pretrained=True).eval().to("cuda")

tests/py/api/test_embed_engines.py

@@ -7,6 +7,7 @@
 from typing import Dict
 from utils import cosine_similarity, COSINE_THRESHOLD
 
+
 class TestModelToEngineToModel(unittest.TestCase):
     def test_resnet50(self):
         self.model = models.resnet50(pretrained=True).eval().to("cuda")
@@ -26,13 +27,20 @@ def test_resnet50(self):
         }
 
         self.scripted_model = torch.jit.script(self.model)
-        trt_engine = torchtrt.ts.convert_method_to_trt_engine(self.scripted_model, "forward", **compile_spec)
+        trt_engine = torchtrt.ts.convert_method_to_trt_engine(
+            self.scripted_model, "forward", **compile_spec
+        )
         trt_mod = torchtrt.ts.embed_engine_in_new_module(trt_engine)
         cos_sim = cosine_similarity(self.model(self.input), trt_mod(self.input))
-        self.assertTrue(cos_sim > COSINE_THRESHOLD, msg=f"Resnet50 TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}")
+        self.assertTrue(
+            cos_sim > COSINE_THRESHOLD,
+            msg=f"Resnet50 TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}",
+        )
 
     def test_efficientnet_b0(self):
-        self.model = timm.create_model("efficientnet_b0", pretrained=True).eval().to("cuda")
+        self.model = (
+            timm.create_model("efficientnet_b0", pretrained=True).eval().to("cuda")
+        )
         self.input = torch.randn((1, 3, 224, 224)).to("cuda")
 
         compile_spec = {
@@ -49,11 +57,17 @@ def test_efficientnet_b0(self):
         }
 
         self.scripted_model = torch.jit.script(self.model)
-        trt_engine = torchtrt.ts.convert_method_to_trt_engine(self.scripted_model, "forward", **compile_spec)
+        trt_engine = torchtrt.ts.convert_method_to_trt_engine(
+            self.scripted_model, "forward", **compile_spec
+        )
         trt_mod = torchtrt.ts.embed_engine_in_new_module(trt_engine)
 
         cos_sim = cosine_similarity(self.model(self.input), trt_mod(self.input))
-        self.assertTrue(cos_sim > COSINE_THRESHOLD, msg=f"EfficientNet-B0 TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}")
+        self.assertTrue(
+            cos_sim > COSINE_THRESHOLD,
+            msg=f"EfficientNet-B0 TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}",
+        )
+
 
 if __name__ == "__main__":
     unittest.main()

tests/py/api/test_module_fallback.py

@@ -6,6 +6,7 @@
 from typing import Dict
 from utils import cosine_similarity, COSINE_THRESHOLD
 
+
 class TestModuleFallback(unittest.TestCase):
     def test_fallback_resnet18(self):
         self.model = models.resnet18(pretrained=True).eval().to("cuda")
@@ -25,7 +26,10 @@ def test_fallback_resnet18(self):
         }
         trt_mod = torchtrt.compile(self.model, **compile_spec)
         cos_sim = cosine_similarity(self.model(self.input), trt_mod(self.input))
-        self.assertTrue(cos_sim > COSINE_THRESHOLD, msg=f"Resnet18 TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}")
+        self.assertTrue(
+            cos_sim > COSINE_THRESHOLD,
+            msg=f"Resnet18 TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}",
+        )
 
     def test_fallback_mobilenet_v2(self):
         self.model = models.mobilenet_v2(pretrained=True).eval().to("cuda")
@@ -41,12 +45,17 @@ def test_fallback_mobilenet_v2(self):
                 "gpu_id": 0,
             },
             "enabled_precisions": {torch.float},
-            "torch_executed_modules": ["torchvision.models.mobilenetv2.ConvBNActivation"],
+            "torch_executed_modules": [
+                "torchvision.models.mobilenetv2.ConvBNActivation"
+            ],
             "min_block_size": 5,
         }
         trt_mod = torchtrt.compile(self.model, **compile_spec)
         cos_sim = cosine_similarity(self.model(self.input), trt_mod(self.input))
-        self.assertTrue(cos_sim > COSINE_THRESHOLD, msg=f"Mobilenet V2 TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}")
+        self.assertTrue(
+            cos_sim > COSINE_THRESHOLD,
+            msg=f"Mobilenet V2 TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}",
+        )
 
 
 if __name__ == "__main__":

tests/py/api/test_operator_fallback.py

@@ -6,6 +6,7 @@
 from typing import Dict
 from utils import cosine_similarity, COSINE_THRESHOLD
 
+
 class TestFallbackModels(unittest.TestCase):
     def test_fallback_resnet18(self):
         self.model = models.resnet18(pretrained=True).eval().to("cuda")
@@ -25,7 +26,10 @@ def test_fallback_resnet18(self):
         }
         trt_mod = torchtrt.compile(self.model, **compile_spec)
         cos_sim = cosine_similarity(self.model(self.input), trt_mod(self.input))
-        self.assertTrue(cos_sim > COSINE_THRESHOLD, msg=f"Resnet18 TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}")
+        self.assertTrue(
+            cos_sim > COSINE_THRESHOLD,
+            msg=f"Resnet18 TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}",
+        )
 
     def test_fallback_mobilenet_v2(self):
         self.model = models.mobilenet_v2(pretrained=True).eval().to("cuda")
@@ -45,7 +49,10 @@ def test_fallback_mobilenet_v2(self):
         }
         trt_mod = torchtrt.compile(self.model, **compile_spec)
         cos_sim = cosine_similarity(self.model(self.input), trt_mod(self.input))
-        self.assertTrue(cos_sim > COSINE_THRESHOLD, msg=f"Mobilenet V2 TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}")
+        self.assertTrue(
+            cos_sim > COSINE_THRESHOLD,
+            msg=f"Mobilenet V2 TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}",
+        )
 
 
 if __name__ == "__main__":