make torchvision as optional dependency for torchscript tests (#3872)

Author: lanluo-nvidia

.github/workflows/build_linux.yml

@@ -284,7 +284,7 @@ jobs:
                   TRT_VERSION=$(cat dev_dep_versions.yml | grep __tensorrt_version__ | sed 's/__tensorrt_version__: //g' | sed 's/"//g')
                   bazel build //:libtorchtrt --compilation_mode opt --config=linux
                   cp bazel-bin/libtorchtrt.tar.gz \
-                  release/tarball/libtorchtrt-${BUILD_VERSION}-tensorrt${TRT_VERSION}-cuda${CU_VERSION:2}-libtorch${PYTORCH_VERSION}-x86_64-linux.tar.gz
+                  release/tarball/libtorchtrt-${BUILD_VERSION}-tensorrt${TRT_VERSION}-cuda${CU_VERSION:2}-libtorch${PYTORCH_VERSION}-${{inputs.architecture}}-linux.tar.gz
                   return 0
               else
                   if [[ ${{ inputs.is-release-wheel }} == true ]]; then

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")