fix: skip tests that require torchvision operators

Author: narendasan

tests/py/ts/api/test_e2e_behavior.py

@@ -1,22 +1,25 @@
 import copy
-import importlib.util
 import unittest
 from typing import Dict
 
 import torch
 import torch_tensorrt as torchtrt
 from utils import same_output_format
 
-if importlib.util.find_spec("torchvision"):
+try:
     import torchvision.models as models
 
+    HAS_TORCHVISION = True
+except (ImportError, RuntimeError):
+    HAS_TORCHVISION = False
+
 
 @unittest.skipIf(
     torchtrt.ENABLED_FEATURES.tensorrt_rtx,
     "aten::adaptive_avg_pool2d is implemented via plugins which is not supported for tensorrt_rtx",
 )
 @unittest.skipIf(
-    not importlib.util.find_spec("torchvision"), "torchvision not installed"
+    not HAS_TORCHVISION, "torchvision not available"
 )
 class TestInputTypeDefaultsFP32Model(unittest.TestCase):
 
@@ -68,7 +71,7 @@ class TestInputTypeDefaultsFP16Model(unittest.TestCase):
         "aten::adaptive_avg_pool2d is implemented via plugins which is not supported for tensorrt_rtx",
     )
     @unittest.skipIf(
-        not importlib.util.find_spec("torchvision"), "torchvision not installed"
+        not HAS_TORCHVISION, "torchvision not available"
     )
     def test_input_use_default_fp16(self):
         self.model = models.resnet18(pretrained=True).eval().to("cuda")
@@ -89,7 +92,7 @@ def test_input_use_default_fp16(self):
         "aten::adaptive_avg_pool2d is implemented via plugins which is not supported for tensorrt_rtx",
     )
     @unittest.skipIf(
-        not importlib.util.find_spec("torchvision"), "torchvision not installed"
+        not HAS_TORCHVISION, "torchvision not available"
     )
     def test_input_use_default_fp16_without_fp16_enabled(self):
         self.model = models.resnet18(pretrained=True).eval().to("cuda")
@@ -108,7 +111,7 @@ def test_input_use_default_fp16_without_fp16_enabled(self):
         "aten::adaptive_avg_pool2d is implemented via plugins which is not supported for tensorrt_rtx",
     )
     @unittest.skipIf(
-        not importlib.util.find_spec("torchvision"), "torchvision not installed"
+        not HAS_TORCHVISION, "torchvision not available"
     )
     def test_input_respect_user_setting_fp16_weights_fp32_in(self):
         self.model = models.resnet18(pretrained=True).eval().to("cuda")
@@ -130,7 +133,7 @@ def test_input_respect_user_setting_fp16_weights_fp32_in(self):
         "aten::adaptive_avg_pool2d is implemented via plugins which is not supported for tensorrt_rtx",
     )
     @unittest.skipIf(
-        not importlib.util.find_spec("torchvision"), "torchvision not installed"
+        not HAS_TORCHVISION, "torchvision not available"
     )
     def test_input_respect_user_setting_fp16_weights_fp32_in_non_constuctor(self):
         self.model = models.resnet18(pretrained=True).eval().to("cuda")

tests/py/ts/api/test_embed_engines.py

@@ -1,20 +1,23 @@
 import copy
-import importlib
 import unittest
 from typing import Dict
 
 import torch
 import torch_tensorrt as torchtrt
 from utils import COSINE_THRESHOLD, cosine_similarity
 
-if importlib.util.find_spec("torchvision"):
+try:
     import timm
     import torchvision.models as models
 
+    HAS_TORCHVISION = True
+except (ImportError, RuntimeError):
+    HAS_TORCHVISION = False
+
 
 class TestModelToEngineToModel(unittest.TestCase):
     @unittest.skipIf(
-        not importlib.util.find_spec("torchvision"), "torchvision not installed"
+        not HAS_TORCHVISION, "torchvision not available"
     )
     @unittest.skipIf(
         torchtrt.ENABLED_FEATURES.tensorrt_rtx,
@@ -49,9 +52,8 @@ def test_resnet50(self):
         )
 
     @unittest.skipIf(
-        not importlib.util.find_spec("timm")
-        or not importlib.util.find_spec("torchvision"),
-        "timm or torchvision not installed",
+        not HAS_TORCHVISION,
+        "timm or torchvision not available",
     )
     @unittest.skipIf(
         torchtrt.ENABLED_FEATURES.tensorrt_rtx,

tests/py/ts/api/test_module_fallback.py

@@ -1,20 +1,23 @@
-import importlib.util
 import unittest
 
 import torch
 import torch_tensorrt as torchtrt
 from utils import COSINE_THRESHOLD, cosine_similarity
 
-if importlib.util.find_spec("torchvision"):
+try:
     import torchvision.models as models
 
+    HAS_TORCHVISION = True
+except (ImportError, RuntimeError):
+    HAS_TORCHVISION = False
+
 
 @unittest.skipIf(
     torchtrt.ENABLED_FEATURES.tensorrt_rtx,
     "aten::adaptive_avg_pool2d is implemented via plugins which is not supported for tensorrt_rtx",
 )
 @unittest.skipIf(
-    not importlib.util.find_spec("torchvision"), "torchvision not installed"
+    not HAS_TORCHVISION, "torchvision not available"
 )
 class TestModuleFallback(unittest.TestCase):
     def test_fallback_resnet18(self):

tests/py/ts/api/test_operator_fallback.py

@@ -1,20 +1,23 @@
-import importlib.util
 import unittest
 
 import torch
 import torch_tensorrt as torchtrt
 from utils import COSINE_THRESHOLD, cosine_similarity
 
-if importlib.util.find_spec("torchvision"):
+try:
     import torchvision.models as models
 
+    HAS_TORCHVISION = True
+except (ImportError, RuntimeError):
+    HAS_TORCHVISION = False
+
 
 @unittest.skipIf(
     torchtrt.ENABLED_FEATURES.tensorrt_rtx,
     "aten::adaptive_avg_pool2d is implemented via plugins which is not supported for tensorrt_rtx",
 )
 @unittest.skipIf(
-    not importlib.util.find_spec("torchvision"), "torchvision not installed"
+    not HAS_TORCHVISION, "torchvision not available"
 )
 class TestFallbackModels(unittest.TestCase):
     def test_fallback_resnet18(self):

tests/py/ts/api/test_ts_backend.py

@@ -1,18 +1,21 @@
 import copy
-import importlib.util
 import unittest
 from typing import Dict
 
 import torch
 import torch_tensorrt as torchtrt
 from utils import COSINE_THRESHOLD, cosine_similarity
 
-if importlib.util.find_spec("torchvision"):
+try:
     import torchvision.models as models
 
+    HAS_TORCHVISION = True
+except (ImportError, RuntimeError):
+    HAS_TORCHVISION = False
+
 
 @unittest.skipIf(
-    not importlib.util.find_spec("torchvision"), "torchvision not installed"
+    not HAS_TORCHVISION, "torchvision not available"
 )
 class TestCompile(unittest.TestCase):
     def test_compile_traced(self):
@@ -129,7 +132,7 @@ def test_default_device(self):
     "aten::adaptive_avg_pool2d is implemented via plugins which is not supported for tensorrt_rtx",
 )
 @unittest.skipIf(
-    not importlib.util.find_spec("torchvision"), "torchvision not installed"
+    not HAS_TORCHVISION, "torchvision not available"
 )
 class TestCheckMethodOpSupport(unittest.TestCase):
     def test_check_support(self):
@@ -138,7 +141,9 @@ def test_check_support(self):
 
         self.assertTrue(torchtrt.ts.check_method_op_support(self.module, "forward"))
 
-
+@unittest.skipIf(
+    not HAS_TORCHVISION, "torchvision not available"
+)
 class TestModuleIdentification(unittest.TestCase):
     def test_module_type(self):
         nn_module = models.alexnet(pretrained=True).eval().to("cuda")