remove torchvision dependency from build, optional for test

.github/scripts/install-torch-tensorrt.sh

@@ -1,7 +1,8 @@
 #set -exou pipefail
 set -x
 
-TORCH_TORCHVISION=$(grep "^torch" ${PWD}/py/requirements.txt)
+TORCH=$(grep "^torch>" ${PWD}/py/requirements.txt)
+TORCHVISION=$(grep "^torchvision" ${PWD}/py/requirements.txt)
 INDEX_URL=https://download.pytorch.org/whl/${CHANNEL}/${CU_VERSION}
 PLATFORM=$(python -c "import sys; print(sys.platform)")
 
@@ -12,8 +13,10 @@ if [[ $(uname -m) == "aarch64" ]]; then
 fi
 
 # Install all the dependencies required for Torch-TensorRT
-pip install --pre ${TORCH_TORCHVISION} --index-url ${INDEX_URL}
 pip install --pre -r ${PWD}/tests/py/requirements.txt
+pip install --force-reinstall --pre ${TORCH} --index-url ${INDEX_URL}
+pip install --force-reinstall --pre ${TORCHVISION} --index-url ${INDEX_URL}
+
 
 # Install Torch-TensorRT
 if [[ ${PLATFORM} == win32 ]]; then

packaging/pre_build_script.sh

@@ -42,14 +42,14 @@ curl -L https://github.com/bazelbuild/bazelisk/releases/download/v1.26.0/bazelis
 pip uninstall -y torch torchvision
 
 if [[ ${IS_JETPACK} == true ]]; then
-    # install torch 2.7 torchvision 0.22.0 for jp6.2
-    pip install torch==2.7.0 torchvision==0.22.0  --index-url=https://pypi.jetson-ai-lab.dev/jp6/cu126/
+    # install torch 2.7 for jp6.2
+    pip install torch==2.7.0 --index-url=https://pypi.jetson-ai-lab.dev/jp6/cu126/
 else
-    TORCH_TORCHVISION=$(grep "^torch" py/requirements.txt)
+    TORCH=$(grep "^torch>" py/requirements.txt)
     INDEX_URL=https://download.pytorch.org/whl/${CHANNEL}/${CU_VERSION}
 
     # Install all the dependencies required for Torch-TensorRT
-    pip install --force-reinstall --pre ${TORCH_TORCHVISION} --index-url ${INDEX_URL}
+    pip install --force-reinstall --pre ${TORCH} --index-url ${INDEX_URL}
 fi
 
 export TORCH_BUILD_NUMBER=$(python -c "import torch, urllib.parse as ul; print(ul.quote_plus(torch.__version__))")

packaging/pre_build_script_windows.sh

@@ -25,12 +25,12 @@ if [[ "${CU_VERSION::4}" < "cu12" ]]; then
          pyproject.toml
 fi
 
-TORCH_TORCHVISION=$(grep "^torch" py/requirements.txt)
+TORCH=$(grep "^torch>" py/requirements.txt)
 INDEX_URL=https://download.pytorch.org/whl/${CHANNEL}/${CU_VERSION}
 
 # Install all the dependencies required for Torch-TensorRT
 pip uninstall -y torch torchvision
-pip install --force-reinstall --pre ${TORCH_TORCHVISION} --index-url ${INDEX_URL}
+pip install --force-reinstall --pre ${TORCH} --index-url ${INDEX_URL}
 
 export CUDA_HOME="$(echo ${CUDA_PATH} | sed -e 's#\\#\/#g')"
 export TORCH_INSTALL_PATH="$(python -c "import torch, os; print(os.path.dirname(torch.__file__))" | sed -e 's#\\#\/#g')"

tests/modules/hub.py

@@ -1,10 +1,13 @@
+import importlib
 import json
 import os
 
 import custom_models as cm
-import timm
 import torch
-import torchvision.models as models
+
+if importlib.util.find_spec("torchvision"):
+    import timm
+    import torchvision.models as models
 
 torch.hub._validate_not_a_forked_repo = lambda a, b, c: True
 
@@ -19,53 +22,7 @@
 # Downloads all model files again if manifest file is not present
 MANIFEST_FILE = "model_manifest.json"
 
-models = {
-    "alexnet": {"model": models.alexnet(pretrained=True), "path": "both"},
-    "vgg16": {"model": models.vgg16(pretrained=True), "path": "both"},
-    "squeezenet": {"model": models.squeezenet1_0(pretrained=True), "path": "both"},
-    "densenet": {"model": models.densenet161(pretrained=True), "path": "both"},
-    "inception_v3": {"model": models.inception_v3(pretrained=True), "path": "both"},
-    "shufflenet": {"model": models.shufflenet_v2_x1_0(pretrained=True), "path": "both"},
-    "mobilenet_v2": {"model": models.mobilenet_v2(pretrained=True), "path": "both"},
-    "resnext50_32x4d": {
-        "model": models.resnext50_32x4d(pretrained=True),
-        "path": "both",
-    },
-    "wideresnet50_2": {
-        "model": models.wide_resnet50_2(pretrained=True),
-        "path": "both",
-    },
-    "mnasnet": {"model": models.mnasnet1_0(pretrained=True), "path": "both"},
-    "resnet18": {
-        "model": torch.hub.load("pytorch/vision:v0.9.0", "resnet18", pretrained=True),
-        "path": "both",
-    },
-    "resnet50": {
-        "model": torch.hub.load("pytorch/vision:v0.9.0", "resnet50", pretrained=True),
-        "path": "both",
-    },
-    "efficientnet_b0": {
-        "model": timm.create_model("efficientnet_b0", pretrained=True),
-        "path": "script",
-    },
-    "vit": {
-        "model": timm.create_model("vit_base_patch16_224", pretrained=True),
-        "path": "script",
-    },
-    "pooling": {"model": cm.Pool(), "path": "trace"},
-    "module_fallback": {"model": cm.ModuleFallbackMain(), "path": "script"},
-    "loop_fallback_eval": {"model": cm.LoopFallbackEval(), "path": "script"},
-    "loop_fallback_no_eval": {"model": cm.LoopFallbackNoEval(), "path": "script"},
-    "conditional": {"model": cm.FallbackIf(), "path": "script"},
-    "inplace_op_if": {"model": cm.FallbackInplaceOPIf(), "path": "script"},
-    "standard_tensor_input": {"model": cm.StandardTensorInput(), "path": "script"},
-    "tuple_input": {"model": cm.TupleInput(), "path": "script"},
-    "list_input": {"model": cm.ListInput(), "path": "script"},
-    "tuple_input_output": {"model": cm.TupleInputOutput(), "path": "script"},
-    "list_input_output": {"model": cm.ListInputOutput(), "path": "script"},
-    "list_input_tuple_output": {"model": cm.ListInputTupleOutput(), "path": "script"},
-    # "bert_base_uncased": {"model": cm.BertModule(), "path": "trace"},
-}
+models = {}
 
 
 def get(n, m, manifest):
@@ -120,6 +77,68 @@ def download_models(version_matches, manifest):
 
 
 def main():
+    if not importlib.util.find_spec("torchvision"):
+        print(f"torchvision is not installed, skip models download")
+        return
+
+    models = {
+        "alexnet": {"model": models.alexnet(pretrained=True), "path": "both"},
+        "vgg16": {"model": models.vgg16(pretrained=True), "path": "both"},
+        "squeezenet": {"model": models.squeezenet1_0(pretrained=True), "path": "both"},
+        "densenet": {"model": models.densenet161(pretrained=True), "path": "both"},
+        "inception_v3": {"model": models.inception_v3(pretrained=True), "path": "both"},
+        "shufflenet": {
+            "model": models.shufflenet_v2_x1_0(pretrained=True),
+            "path": "both",
+        },
+        "mobilenet_v2": {"model": models.mobilenet_v2(pretrained=True), "path": "both"},
+        "resnext50_32x4d": {
+            "model": models.resnext50_32x4d(pretrained=True),
+            "path": "both",
+        },
+        "wideresnet50_2": {
+            "model": models.wide_resnet50_2(pretrained=True),
+            "path": "both",
+        },
+        "mnasnet": {"model": models.mnasnet1_0(pretrained=True), "path": "both"},
+        "resnet18": {
+            "model": torch.hub.load(
+                "pytorch/vision:v0.9.0", "resnet18", pretrained=True
+            ),
+            "path": "both",
+        },
+        "resnet50": {
+            "model": torch.hub.load(
+                "pytorch/vision:v0.9.0", "resnet50", pretrained=True
+            ),
+            "path": "both",
+        },
+        "efficientnet_b0": {
+            "model": timm.create_model("efficientnet_b0", pretrained=True),
+            "path": "script",
+        },
+        "vit": {
+            "model": timm.create_model("vit_base_patch16_224", pretrained=True),
+            "path": "script",
+        },
+        "pooling": {"model": cm.Pool(), "path": "trace"},
+        "module_fallback": {"model": cm.ModuleFallbackMain(), "path": "script"},
+        "loop_fallback_eval": {"model": cm.LoopFallbackEval(), "path": "script"},
+        "loop_fallback_no_eval": {"model": cm.LoopFallbackNoEval(), "path": "script"},
+        "conditional": {"model": cm.FallbackIf(), "path": "script"},
+        "inplace_op_if": {"model": cm.FallbackInplaceOPIf(), "path": "script"},
+        "standard_tensor_input": {"model": cm.StandardTensorInput(), "path": "script"},
+        "tuple_input": {"model": cm.TupleInput(), "path": "script"},
+        "list_input": {"model": cm.ListInput(), "path": "script"},
+        "tuple_input_output": {"model": cm.TupleInputOutput(), "path": "script"},
+        "list_input_output": {"model": cm.ListInputOutput(), "path": "script"},
+        "list_input_tuple_output": {
+            "model": cm.ListInputTupleOutput(),
+            "path": "script",
+        },
+        # "bert_base_uncased": {"model": cm.BertModule(), "path": "trace"},
+    }
+
     manifest = None
     version_matches = False
     manifest_exists = False

tests/py/core/test_classes.py

@@ -2,14 +2,12 @@
 import unittest
 from typing import Dict
 
+import tensorrt as trt
 import torch
 import torch_tensorrt
 import torch_tensorrt as torchtrt
-import torchvision.models as models
 from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import TorchTensorRTModule
 
-import tensorrt as trt
-
 
 class TestDevice(unittest.TestCase):
     def test_from_string_constructor(self):

tests/py/dynamo/models/test_dyn_models.py

@@ -1,9 +1,8 @@
 # type: ignore
-
+import importlib
 import unittest
 
 import pytest
-import timm
 import torch
 import torch_tensorrt as torchtrt
 from torch_tensorrt.dynamo.utils import COSINE_THRESHOLD, cosine_similarity
@@ -175,6 +174,9 @@ def forward(self, x):
     )
 
 
+@unittest.skipIf(
+    not importlib.util.find_spec("torchvision"), "torchvision not installed"
+)
 @pytest.mark.unit
 def test_resnet_dynamic(ir):
     """

tests/py/dynamo/models/test_engine_cache.py

@@ -1,4 +1,5 @@
 # type: ignore
+import importlib
 import os
 import shutil
 import unittest
@@ -7,7 +8,6 @@
 import pytest
 import torch
 import torch_tensorrt as torch_trt
-import torchvision.models as models
 from torch.testing._internal.common_utils import TestCase
 from torch_tensorrt.dynamo._defaults import TIMING_CACHE_PATH
 from torch_tensorrt.dynamo._engine_cache import BaseEngineCache
@@ -16,6 +16,9 @@
 
 assertions = unittest.TestCase()
 
+if importlib.util.find_spec("torchvision"):
+    import torchvision.models as models
+
 
 class MyEngineCache(BaseEngineCache):
     def __init__(
@@ -57,6 +60,9 @@ def load(self, hash: str, prefix: str = "blob") -> Optional[bytes]:
 
 
 class TestHashFunction(TestCase):
+    @unittest.skipIf(
+        not importlib.util.find_spec("torchvision"), "torchvision not installed"
+    )
     def test_reexport_is_equal(self):
         pyt_model = models.resnet18(pretrained=True).eval().to("cuda")
         example_inputs = (torch.randn((100, 3, 224, 224)).to("cuda"),)
@@ -94,6 +100,9 @@ def test_reexport_is_equal(self):
 
         self.assertEqual(hash1, hash2)
 
+    @unittest.skipIf(
+        not importlib.util.find_spec("torchvision"), "torchvision not installed"
+    )
     def test_input_shape_change_is_not_equal(self):
         pyt_model = models.resnet18(pretrained=True).eval().to("cuda")
         example_inputs = (torch.randn((100, 3, 224, 224)).to("cuda"),)
@@ -131,6 +140,9 @@ def test_input_shape_change_is_not_equal(self):
 
         self.assertNotEqual(hash1, hash2)
 
+    @unittest.skipIf(
+        not importlib.util.find_spec("torchvision"), "torchvision not installed"
+    )
     def test_engine_settings_is_not_equal(self):
         pyt_model = models.resnet18(pretrained=True).eval().to("cuda")
         example_inputs = (torch.randn((100, 3, 224, 224)).to("cuda"),)
@@ -177,6 +189,9 @@ def test_engine_settings_is_not_equal(self):
 
 class TestEngineCache(TestCase):
     @pytest.mark.xfail
+    @unittest.skipIf(
+        not importlib.util.find_spec("torchvision"), "torchvision not installed"
+    )
     def test_dynamo_compile_with_default_disk_engine_cache(self):
         model = models.resnet18(pretrained=True).eval().to("cuda")
         example_inputs = (torch.randn((100, 3, 224, 224)).to("cuda"),)
@@ -254,6 +269,9 @@ def remove_timing_cache(path=TIMING_CACHE_PATH):
         not torch_trt.ENABLED_FEATURES.refit,
         "Engine caching requires refit feature that is not supported in Python 3.13 or higher",
     )
+    @unittest.skipIf(
+        not importlib.util.find_spec("torchvision"), "torchvision not installed"
+    )
     def test_dynamo_compile_with_custom_engine_cache(self):
         model = models.resnet18(pretrained=True).eval().to("cuda")
 
@@ -322,6 +340,9 @@ def test_dynamo_compile_with_custom_engine_cache(self):
         not torch_trt.ENABLED_FEATURES.refit,
         "Engine caching requires refit feature that is not supported in Python 3.13 or higher",
     )
+    @unittest.skipIf(
+        not importlib.util.find_spec("torchvision"), "torchvision not installed"
+    )
     def test_dynamo_compile_change_input_shape(self):
         """Runs compilation 3 times, the cache should miss each time"""
         model = models.resnet18(pretrained=True).eval().to("cuda")
@@ -358,6 +379,9 @@ def test_dynamo_compile_change_input_shape(self):
         not torch_trt.ENABLED_FEATURES.refit,
         "Engine caching requires refit feature that is not supported in Python 3.13 or higher",
     )
+    @unittest.skipIf(
+        not importlib.util.find_spec("torchvision"), "torchvision not installed"
+    )
     @pytest.mark.xfail
     def test_torch_compile_with_default_disk_engine_cache(self):
         # Custom Engine Cache
@@ -430,6 +454,9 @@ def remove_timing_cache(path=TIMING_CACHE_PATH):
             msg=f"Engine caching didn't speed up the compilation. Time taken without engine caching: {times[0]} ms, time taken with engine caching: {times[2]} ms",
         )
 
+    @unittest.skipIf(
+        not importlib.util.find_spec("torchvision"), "torchvision not installed"
+    )
     def test_torch_compile_with_custom_engine_cache(self):
         # Custom Engine Cache
         model = models.resnet18(pretrained=True).eval().to("cuda")
@@ -501,6 +528,9 @@ def test_torch_compile_with_custom_engine_cache(self):
         not torch_trt.ENABLED_FEATURES.refit,
         "Engine caching requires refit feature that is not supported in Python 3.13 or higher",
     )
+    @unittest.skipIf(
+        not importlib.util.find_spec("torchvision"), "torchvision not installed"
+    )
     def test_torch_trt_compile_change_input_shape(self):
         # Custom Engine Cache
         model = models.resnet18(pretrained=True).eval().to("cuda")
@@ -631,6 +661,9 @@ def forward(self, c, d):
         not torch_trt.ENABLED_FEATURES.refit,
         "Engine caching requires refit feature that is not supported in Python 3.13 or higher",
     )
+    @unittest.skipIf(
+        not importlib.util.find_spec("torchvision"), "torchvision not installed"
+    )
     def test_caching_small_model(self):
         from torch_tensorrt.dynamo._refit import refit_module_weights
 

tests/py/dynamo/models/test_export_kwargs_serde.py

@@ -4,11 +4,9 @@
 import unittest
 
 import pytest
-import timm
 import torch
 import torch.nn.functional as F
 import torch_tensorrt as torchtrt
-import torchvision.models as models
 from torch import nn
 from torch_tensorrt.dynamo._compiler import (
     convert_exported_program_to_serialized_trt_engine,