make torchvision as optional dependency for torchscript tests

tests/py/ts/api/test_classes.py

@@ -4,7 +4,6 @@
 
 import torch
 import torch_tensorrt as torchtrt
-import torchvision.models as models
 from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import TorchTensorRTModule
 
 

tests/py/ts/api/test_collections.py

@@ -5,7 +5,6 @@
 
 import torch
 import torch_tensorrt as torchtrt
-import torchvision.models as models
 from utils import COSINE_THRESHOLD, cosine_similarity
 
 

tests/py/ts/api/test_e2e_behavior.py

@@ -1,18 +1,25 @@
 import copy
+import importlib.util
 import unittest
 from typing import Dict
 
 import torch
 import torch_tensorrt as torchtrt
-import torchvision.models as models
 from utils import same_output_format
 
+if importlib.util.find_spec("torchvision"):
+    import torchvision.models as models
+
 
 @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"
+)
 class TestInputTypeDefaultsFP32Model(unittest.TestCase):
+
     def test_input_use_default_fp32(self):
         self.model = models.resnet18(pretrained=True).eval().to("cuda")
         self.input = torch.randn((1, 3, 224, 224)).to("cuda")
@@ -60,6 +67,9 @@ class TestInputTypeDefaultsFP16Model(unittest.TestCase):
         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"
+    )
     def test_input_use_default_fp16(self):
         self.model = models.resnet18(pretrained=True).eval().to("cuda")
         self.input = torch.randn((1, 3, 224, 224)).to("cuda")
@@ -78,6 +88,9 @@ def test_input_use_default_fp16(self):
         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"
+    )
     def test_input_use_default_fp16_without_fp16_enabled(self):
         self.model = models.resnet18(pretrained=True).eval().to("cuda")
         self.input = torch.randn((1, 3, 224, 224)).to("cuda")
@@ -94,6 +107,9 @@ def test_input_use_default_fp16_without_fp16_enabled(self):
         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"
+    )
     def test_input_respect_user_setting_fp16_weights_fp32_in(self):
         self.model = models.resnet18(pretrained=True).eval().to("cuda")
         self.input = torch.randn((1, 3, 224, 224)).to("cuda")
@@ -113,6 +129,9 @@ def test_input_respect_user_setting_fp16_weights_fp32_in(self):
         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"
+    )
     def test_input_respect_user_setting_fp16_weights_fp32_in_non_constuctor(self):
         self.model = models.resnet18(pretrained=True).eval().to("cuda")
         self.input = torch.randn((1, 3, 224, 224)).to("cuda")

tests/py/ts/api/test_logging.py

@@ -1,10 +1,6 @@
-import copy
 import unittest
-from typing import Dict
 
-import torch
 import torch_tensorrt as torchtrt
-import torchvision.models as models
 
 
 @unittest.skipIf(

tests/py/ts/api/test_module_fallback.py

@@ -1,17 +1,21 @@
-import copy
+import importlib.util
 import unittest
-from typing import Dict
 
 import torch
 import torch_tensorrt as torchtrt
-import torchvision.models as models
 from utils import COSINE_THRESHOLD, cosine_similarity
 
+if importlib.util.find_spec("torchvision"):
+    import torchvision.models as models
+
 
 @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"
+)
 class TestModuleFallback(unittest.TestCase):
     def test_fallback_resnet18(self):
         self.model = models.resnet18(pretrained=True).eval().to("cuda")

tests/py/ts/api/test_operator_fallback.py

@@ -1,17 +1,21 @@
-import copy
+import importlib.util
 import unittest
-from typing import Dict
 
 import torch
 import torch_tensorrt as torchtrt
-import torchvision.models as models
 from utils import COSINE_THRESHOLD, cosine_similarity
 
+if importlib.util.find_spec("torchvision"):
+    import torchvision.models as models
+
 
 @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"
+)
 class TestFallbackModels(unittest.TestCase):
     def test_fallback_resnet18(self):
         self.model = models.resnet18(pretrained=True).eval().to("cuda")

tests/py/ts/api/test_ts_backend.py

@@ -1,13 +1,19 @@
 import copy
+import importlib.util
 import unittest
 from typing import Dict
 
 import torch
 import torch_tensorrt as torchtrt
-import torchvision.models as models
 from utils import COSINE_THRESHOLD, cosine_similarity
 
+if importlib.util.find_spec("torchvision"):
+    import torchvision.models as models
 
+
+@unittest.skipIf(
+    not importlib.util.find_spec("torchvision"), "torchvision not installed"
+)
 class TestCompile(unittest.TestCase):
     def test_compile_traced(self):
         self.model = models.vgg16(pretrained=True).eval().to("cuda")
@@ -122,6 +128,9 @@ def test_default_device(self):
     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"
+)
 class TestCheckMethodOpSupport(unittest.TestCase):
     def test_check_support(self):
         module = models.alexnet(pretrained=True).eval().to("cuda")

tests/py/ts/hw/test_multi_gpu.py

@@ -1,15 +1,21 @@
+import importlib.util
 import unittest
 
 import torch
 import torch_tensorrt as torchtrt
-import torchvision.models as models
 from model_test_case import ModelTestCase
 
+if importlib.util.find_spec("torchvision"):
+    import torchvision.models as models
+
 
 @unittest.skipIf(
     not torchtrt.ENABLED_FEATURES.torchscript_frontend,
     "TorchScript Frontend is not available",
 )
+@unittest.skipIf(
+    not importlib.util.find_spec("torchvision"), "torchvision not installed"
+)
 class TestMultiGpuSwitching(ModelTestCase):
     def setUp(self):
         if torch.cuda.device_count() < 2:
@@ -73,6 +79,9 @@ def test_compile_script(self):
     not torchtrt.ENABLED_FEATURES.torchscript_frontend,
     "TorchScript Frontend is not available",
 )
+@unittest.skipIf(
+    not importlib.util.find_spec("torchvision"), "torchvision not installed"
+)
 class TestMultiGpuSerializeDeserializeSwitching(ModelTestCase):
     def setUp(self):
         if torch.cuda.device_count() < 2:
@@ -148,9 +157,10 @@ def test_suite():
     return suite
 
 
-suite = test_suite()
+if importlib.util.find_spec("torchvision"):
+    suite = test_suite()
 
-runner = unittest.TextTestRunner()
-result = runner.run(suite)
+    runner = unittest.TextTestRunner()
+    result = runner.run(suite)
 
-exit(int(not result.wasSuccessful()))
+    exit(int(not result.wasSuccessful()))

tests/py/ts/integrations/test_to_backend_api.py

@@ -1,11 +1,14 @@
 # type: ignore
+import importlib.util
 import unittest
 
 import torch
 import torch_tensorrt as torchtrt
-import torchvision.models as models
 from utils import COSINE_THRESHOLD, cosine_similarity
 
+if importlib.util.find_spec("torchvision"):
+    import torchvision.models as models
+
 
 @unittest.skipIf(
     not torchtrt.ENABLED_FEATURES.torchscript_frontend,
@@ -15,6 +18,9 @@
     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"
+)
 class TestToBackendLowering(unittest.TestCase):
     def setUp(self):
         self.input = torch.randn((1, 3, 300, 300)).to("cuda")

tests/py/ts/integrations/test_trt_intercompatibility.py

@@ -1,10 +1,13 @@
+import importlib.util
 import unittest
 
 import torch
 import torch_tensorrt as torchtrt
-import torchvision.models as models
 from utils import COSINE_THRESHOLD, cosine_similarity
 
+if importlib.util.find_spec("torchvision"):
+    import torchvision.models as models
+
 
 @unittest.skipIf(
     not torchtrt.ENABLED_FEATURES.torchscript_frontend,
@@ -14,6 +17,9 @@
     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"
+)
 class TestPyTorchToTRTEngine(unittest.TestCase):
     def test_pt_to_trt(self):
         self.model = models.resnet18(pretrained=True).eval().to("cuda:0")