diff --git a/.github/scripts/install-torch-tensorrt.sh b/.github/scripts/install-torch-tensorrt.sh old mode 100644 new mode 100755 index fd88207eb4..109b9565f3 --- a/.github/scripts/install-torch-tensorrt.sh +++ b/.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 diff --git a/packaging/pre_build_script.sh b/packaging/pre_build_script.sh index f7cbb4919e..707edcd718 100755 --- a/packaging/pre_build_script.sh +++ b/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__))") diff --git a/packaging/pre_build_script_windows.sh b/packaging/pre_build_script_windows.sh index f73f8875aa..7d089b602a 100644 --- a/packaging/pre_build_script_windows.sh +++ b/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')" diff --git a/tests/modules/hub.py b/tests/modules/hub.py index d87635b435..4cdd3d6350 100644 --- a/tests/modules/hub.py +++ b/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 diff --git a/tests/py/core/test_classes.py b/tests/py/core/test_classes.py index 62abeb6b1a..5f6df00ad8 100644 --- a/tests/py/core/test_classes.py +++ b/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): diff --git a/tests/py/dynamo/models/test_dyn_models.py b/tests/py/dynamo/models/test_dyn_models.py index d5627499f5..052a268a64 100644 --- a/tests/py/dynamo/models/test_dyn_models.py +++ b/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): """ diff --git a/tests/py/dynamo/models/test_engine_cache.py b/tests/py/dynamo/models/test_engine_cache.py index 0bc7c665b3..5529c013cf 100644 --- a/tests/py/dynamo/models/test_engine_cache.py +++ b/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 diff --git a/tests/py/dynamo/models/test_export_kwargs_serde.py b/tests/py/dynamo/models/test_export_kwargs_serde.py index 928d62e7ba..d352829698 100644 --- a/tests/py/dynamo/models/test_export_kwargs_serde.py +++ b/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, diff --git a/tests/py/dynamo/models/test_export_serde.py b/tests/py/dynamo/models/test_export_serde.py index 52e5eefb63..510631a8bc 100644 --- a/tests/py/dynamo/models/test_export_serde.py +++ b/tests/py/dynamo/models/test_export_serde.py @@ -1,3 +1,4 @@ +import importlib import os import tempfile import unittest @@ -5,7 +6,6 @@ import pytest import torch import torch_tensorrt as torchtrt -import torchvision.models as models from torch_tensorrt.dynamo.utils import ( COSINE_THRESHOLD, cosine_similarity, @@ -13,6 +13,8 @@ ) assertions = unittest.TestCase() +if importlib.util.find_spec("torchvision"): + import torchvision.models as models trt_ep_path = os.path.join(tempfile.gettempdir(), "trt.ep") @@ -247,6 +249,10 @@ def forward(self, x): @pytest.mark.unit +@unittest.skipIf( + not importlib.util.find_spec("torchvision"), + "torchvision is not installed", +) def test_resnet18(ir): """ This tests export save and load functionality on Resnet18 model @@ -288,6 +294,10 @@ def test_resnet18(ir): @pytest.mark.unit +@unittest.skipIf( + not importlib.util.find_spec("torchvision"), + "torchvision is not installed", +) def test_resnet18_cpu_offload(ir): """ This tests export save and load functionality on Resnet18 model @@ -335,6 +345,10 @@ def test_resnet18_cpu_offload(ir): @pytest.mark.unit +@unittest.skipIf( + not importlib.util.find_spec("torchvision"), + "torchvision is not installed", +) def test_resnet18_dynamic(ir): """ This tests export save and load functionality on Resnet18 model diff --git a/tests/py/dynamo/models/test_model_refit.py b/tests/py/dynamo/models/test_model_refit.py index d59a9482e2..dc7c21cf66 100644 --- a/tests/py/dynamo/models/test_model_refit.py +++ b/tests/py/dynamo/models/test_model_refit.py @@ -9,7 +9,6 @@ import torch.nn.functional as F import torch_tensorrt as torchtrt import torch_tensorrt as torch_trt -import torchvision.models as models from torch import nn from torch_tensorrt.dynamo import refit_module_weights from torch_tensorrt.dynamo._refit import ( @@ -25,6 +24,9 @@ assertions = unittest.TestCase() +if importlib.util.find_spec("torchvision"): + import torchvision.models as models + @unittest.skipIf( not torch_trt.ENABLED_FEATURES.torch_tensorrt_runtime, @@ -34,6 +36,10 @@ not torch_trt.ENABLED_FEATURES.refit, "Refit feature is not supported in Python 3.13 or higher", ) +@unittest.skipIf( + not importlib.util.find_spec("torchvision"), + "torchvision is not installed", +) @pytest.mark.unit def test_mapping(): model = models.resnet18(pretrained=False).eval().to("cuda") @@ -93,6 +99,10 @@ def test_mapping(): not torch_trt.ENABLED_FEATURES.refit, "Refit feature is not supported in Python 3.13 or higher", ) +@unittest.skipIf( + not importlib.util.find_spec("torchvision"), + "torchvision is not installed", +) @pytest.mark.unit def test_refit_one_engine_with_weightmap(): model = models.resnet18(pretrained=False).eval().to("cuda") @@ -147,6 +157,10 @@ def test_refit_one_engine_with_weightmap(): not torch_trt.ENABLED_FEATURES.refit, "Refit feature is not supported in Python 3.13 or higher", ) +@unittest.skipIf( + not importlib.util.find_spec("torchvision"), + "torchvision is not installed", +) @pytest.mark.unit def test_refit_one_engine_no_map_with_weightmap(): model = models.resnet18(pretrained=False).eval().to("cuda") @@ -202,6 +216,10 @@ def test_refit_one_engine_no_map_with_weightmap(): not torch_trt.ENABLED_FEATURES.refit, "Refit feature is not supported in Python 3.13 or higher", ) +@unittest.skipIf( + not importlib.util.find_spec("torchvision"), + "torchvision is not installed", +) @pytest.mark.unit def test_refit_one_engine_with_wrong_weightmap(): model = models.resnet18(pretrained=False).eval().to("cuda") diff --git a/tests/py/dynamo/models/test_models.py b/tests/py/dynamo/models/test_models.py index b0ebbf5fa4..0a89e24677 100644 --- a/tests/py/dynamo/models/test_models.py +++ b/tests/py/dynamo/models/test_models.py @@ -3,10 +3,8 @@ import unittest import pytest -import timm import torch import torch_tensorrt as torchtrt -import torchvision.models as models from torch_tensorrt.dynamo.utils import ( COSINE_THRESHOLD, cosine_similarity, @@ -15,8 +13,16 @@ assertions = unittest.TestCase() +if importlib.util.find_spec("torchvision"): + import timm + import torchvision.models as models + @pytest.mark.unit +@unittest.skipIf( + not importlib.util.find_spec("torchvision"), + "torchvision is not installed", +) def test_resnet18(ir): model = models.resnet18(pretrained=True).eval().to("cuda") input = torch.randn((1, 3, 224, 224)).to("cuda") @@ -48,6 +54,10 @@ def test_resnet18(ir): @pytest.mark.unit +@unittest.skipIf( + not importlib.util.find_spec("torchvision"), + "torchvision is not installed", +) def test_resnet18_cpu_offload(ir): model = models.resnet18(pretrained=True).eval().to("cuda") input = torch.randn((1, 3, 224, 224)).to("cuda") @@ -85,6 +95,10 @@ def test_resnet18_cpu_offload(ir): @pytest.mark.unit +@unittest.skipIf( + not importlib.util.find_spec("torchvision"), + "torchvision is not installed", +) def test_mobilenet_v2(ir): model = models.mobilenet_v2(pretrained=True).eval().to("cuda") input = torch.randn((1, 3, 224, 224)).to("cuda") @@ -117,6 +131,10 @@ def test_mobilenet_v2(ir): @pytest.mark.unit +@unittest.skipIf( + not importlib.util.find_spec("timm") or not importlib.util.find_spec("torchvision"), + "timm or torchvision not installed", +) def test_efficientnet_b0(ir): model = timm.create_model("efficientnet_b0", pretrained=True).eval().to("cuda") input = torch.randn((1, 3, 224, 224)).to("cuda") @@ -253,6 +271,10 @@ def test_bert_base_uncased_cpu_offload(ir): @pytest.mark.unit +@unittest.skipIf( + not importlib.util.find_spec("torchvision"), + "torchvision is not installed", +) def test_resnet18_half(ir): model = models.resnet18(pretrained=True).eval().to("cuda").half() input = torch.randn((1, 3, 224, 224)).to("cuda").half() diff --git a/tests/py/dynamo/models/test_models_export.py b/tests/py/dynamo/models/test_models_export.py index 189da962b5..86ff042bb5 100644 --- a/tests/py/dynamo/models/test_models_export.py +++ b/tests/py/dynamo/models/test_models_export.py @@ -5,18 +5,23 @@ from importlib import metadata import pytest -import timm import torch import torch_tensorrt as torchtrt -import torchvision.models as models from torch_tensorrt.dynamo.utils import COSINE_THRESHOLD, cosine_similarity from packaging.version import Version +if importlib.util.find_spec("torchvision"): + import timm + import torchvision.models as models + assertions = unittest.TestCase() @pytest.mark.unit +@unittest.skipIf( + not importlib.util.find_spec("torchvision"), "torchvision not installed" +) def test_resnet18(ir): model = models.resnet18(pretrained=True).eval().to("cuda") input = torch.randn((1, 3, 224, 224)).to("cuda") @@ -81,6 +86,10 @@ def test_mobilenet_v2(ir): @pytest.mark.unit +@unittest.skipIf( + not importlib.util.find_spec("timm") or not importlib.util.find_spec("torchvision"), + "timm or torchvision not installed", +) def test_efficientnet_b0(ir): model = timm.create_model("efficientnet_b0", pretrained=True).eval().to("cuda") input = torch.randn((1, 3, 224, 224)).to("cuda") diff --git a/tests/py/dynamo/models/test_reexport.py b/tests/py/dynamo/models/test_reexport.py index 297410ae55..c4c7fb6787 100644 --- a/tests/py/dynamo/models/test_reexport.py +++ b/tests/py/dynamo/models/test_reexport.py @@ -1,3 +1,4 @@ +import importlib import os import tempfile import unittest @@ -6,11 +7,13 @@ import torch import torch.nn as nn import torch_tensorrt as torchtrt -import torchvision.models as models from torch_tensorrt.dynamo.utils import COSINE_THRESHOLD, cosine_similarity assertions = unittest.TestCase() +if importlib.util.find_spec("torchvision"): + import torchvision.models as models + trt_ep_path = os.path.join(tempfile.gettempdir(), "trt.ep") @@ -256,6 +259,10 @@ def forward(self, x): @pytest.mark.unit +@unittest.skipIf( + not importlib.util.find_spec("torchvision"), + "torchvision is not installed", +) def test_resnet18(ir): """ This tests export save and load functionality on Resnet18 model @@ -422,6 +429,10 @@ def forward(self, x): @pytest.mark.unit +@unittest.skipIf( + not importlib.util.find_spec("torchvision"), + "torchvision is not installed", +) def test_resnet18_dynamic(ir): """ This tests export save and load functionality on Resnet18 model with dynamic shapes @@ -492,6 +503,10 @@ def test_resnet18_dynamic(ir): @pytest.mark.unit +@unittest.skipIf( + not importlib.util.find_spec("torchvision"), + "torchvision is not installed", +) def test_resnet18_dynamic_fallback(ir): """ This tests export save and load functionality on Resnet18 model with dynamic shapes and fallback diff --git a/tests/py/dynamo/models/test_weight_stripped_engine.py b/tests/py/dynamo/models/test_weight_stripped_engine.py index 33bf94e711..9110930377 100644 --- a/tests/py/dynamo/models/test_weight_stripped_engine.py +++ b/tests/py/dynamo/models/test_weight_stripped_engine.py @@ -1,3 +1,4 @@ +import importlib import os import pickle import shutil @@ -5,7 +6,6 @@ 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 import convert_exported_program_to_serialized_trt_engine from torch_tensorrt.dynamo._defaults import TIMING_CACHE_PATH @@ -14,12 +14,19 @@ assertions = unittest.TestCase() +if importlib.util.find_spec("torchvision"): + import torchvision.models as models + class TestWeightStrippedEngine(TestCase): @unittest.skipIf( 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 is not installed", + ) def test_three_ways_to_compile(self): pyt_model = models.resnet18(pretrained=True).eval().to("cuda") example_inputs = (torch.randn((100, 3, 224, 224)).to("cuda"),) @@ -65,6 +72,10 @@ def test_three_ways_to_compile(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 is not installed", + ) def test_three_ways_to_compile_weight_stripped_engine(self): pyt_model = models.resnet18(pretrained=True).eval().to("cuda") example_inputs = (torch.randn((100, 3, 224, 224)).to("cuda"),) @@ -101,6 +112,10 @@ def test_three_ways_to_compile_weight_stripped_engine(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 is not installed", + ) def test_weight_stripped_engine_sizes(self): pyt_model = models.resnet18(pretrained=True).eval().to("cuda") example_inputs = (torch.randn((100, 3, 224, 224)).to("cuda"),) @@ -142,6 +157,10 @@ def test_weight_stripped_engine_sizes(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 is not installed", + ) def test_weight_stripped_engine_results(self): pyt_model = models.resnet18(pretrained=True).eval().to("cuda") example_inputs = (torch.randn((100, 3, 224, 224)).to("cuda"),) @@ -207,6 +226,10 @@ def test_weight_stripped_engine_results(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 is not installed", + ) def test_engine_caching_saves_weight_stripped_engine(self): pyt_model = models.resnet18(pretrained=True).eval().to("cuda") example_inputs = (torch.randn((100, 3, 224, 224)).to("cuda"),) @@ -257,6 +280,10 @@ def test_engine_caching_saves_weight_stripped_engine(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 is not installed", + ) def test_dynamo_compile_with_refittable_weight_stripped_engine(self): pyt_model = models.resnet18(pretrained=True).eval().to("cuda") example_inputs = (torch.randn((100, 3, 224, 224)).to("cuda"),) @@ -339,6 +366,10 @@ 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 is not installed", + ) def test_torch_compile_with_refittable_weight_stripped_engine(self): pyt_model = models.resnet18(pretrained=True).eval().to("cuda") @@ -519,6 +550,10 @@ def forward(self, x): 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 is not installed", + ) def test_two_TRTRuntime_in_refitting(self): pyt_model = models.resnet18(pretrained=True).eval().to("cuda") example_inputs = (torch.randn((100, 3, 224, 224)).to("cuda"),) @@ -563,6 +598,10 @@ def test_two_TRTRuntime_in_refitting(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 is not installed", + ) def test_refit_identical_engine_weights(self): pyt_model = models.resnet18(pretrained=True).eval().to("cuda") example_inputs = (torch.randn((100, 3, 224, 224)).to("cuda"),) diff --git a/tests/py/dynamo/runtime/test_002_lazy_engine_init.py b/tests/py/dynamo/runtime/test_002_lazy_engine_init.py index 99df3f2335..ca82797090 100644 --- a/tests/py/dynamo/runtime/test_002_lazy_engine_init.py +++ b/tests/py/dynamo/runtime/test_002_lazy_engine_init.py @@ -1,4 +1,5 @@ # type: ignore +import importlib import os import tempfile import unittest @@ -6,7 +7,6 @@ import torch import torch_tensorrt import torch_tensorrt as torchtrt -import torchvision.models as models from torch.testing._internal.common_utils import TestCase from torch_tensorrt.dynamo import CompilationSettings from torch_tensorrt.dynamo.utils import ( @@ -18,6 +18,9 @@ assertions = unittest.TestCase() +if importlib.util.find_spec("torchvision"): + import torchvision.models as models + def assert_close(outputs, ref_outputs): if type(outputs) not in (list, tuple): @@ -202,6 +205,10 @@ def forward(self, a, b): assert_close(trt_output, model_output) + @unittest.skipIf( + not importlib.util.find_spec("torchvision"), + "torchvision is not installed", + ) def test_lazy_engine_init_py_e2e(self): model = models.resnet18(pretrained=True).eval().to("cuda") input = torch.randn((1, 3, 224, 224)).to("cuda") @@ -238,6 +245,10 @@ def test_lazy_engine_init_py_e2e(self): not torch_tensorrt.ENABLED_FEATURES.torch_tensorrt_runtime, "Torch-TensorRT Runtime is not available", ) + @unittest.skipIf( + not importlib.util.find_spec("torchvision"), + "torchvision is not installed", + ) def test_lazy_engine_init_cpp_e2e(self): model = models.resnet18(pretrained=False).eval().to("cuda") input = torch.randn((1, 3, 224, 224)).to("cuda") @@ -274,6 +285,10 @@ def test_lazy_engine_init_cpp_e2e(self): not torch_tensorrt.ENABLED_FEATURES.torch_tensorrt_runtime, "Torch-TensorRT Runtime is not available", ) + @unittest.skipIf( + not importlib.util.find_spec("torchvision"), + "torchvision is not installed", + ) def test_lazy_engine_init_cpp_serialization(self): model = models.resnet18(pretrained=False).eval().to("cuda") input = torch.randn((1, 3, 224, 224)).to("cuda") diff --git a/tests/py/dynamo/runtime/test_mutable_torchtrt_module.py b/tests/py/dynamo/runtime/test_mutable_torchtrt_module.py index d30f064111..703be00575 100644 --- a/tests/py/dynamo/runtime/test_mutable_torchtrt_module.py +++ b/tests/py/dynamo/runtime/test_mutable_torchtrt_module.py @@ -1,20 +1,22 @@ # type: ignore +import importlib import os import tempfile import unittest import pytest -import timm import torch import torch.nn.functional as F import torch_tensorrt as torch_trt -import torchvision.models as models from torch import nn from torch_tensorrt.dynamo.runtime._MutableTorchTensorRTModule import RefitFlag from torch_tensorrt.dynamo.utils import check_output_equal assertions = unittest.TestCase() +if importlib.util.find_spec("torchvision"): + import torchvision.models as models + @pytest.mark.unit def test_check_output_equal(): @@ -209,6 +211,10 @@ def forward(self, a, b, c=None): not torch_trt.ENABLED_FEATURES.refit, "Refit feature is not supported in Python 3.13 or higher", ) +@unittest.skipIf( + not importlib.util.find_spec("torchvision"), + "torchvision is not installed", +) @pytest.mark.unit def test_resnet18(): torch.manual_seed(0) @@ -249,6 +255,10 @@ def test_resnet18(): not torch_trt.ENABLED_FEATURES.refit, "Refit feature is not supported in Python 3.13 or higher", ) +@unittest.skipIf( + not importlib.util.find_spec("torchvision"), + "torchvision is not installed", +) @pytest.mark.unit def test_save(): torch.manual_seed(0) diff --git a/tests/py/ts/api/test_embed_engines.py b/tests/py/ts/api/test_embed_engines.py index d21e139eca..ed421dfe2c 100644 --- a/tests/py/ts/api/test_embed_engines.py +++ b/tests/py/ts/api/test_embed_engines.py @@ -1,14 +1,21 @@ -import unittest -import torch_tensorrt as torchtrt -import torch -import torchvision.models as models import copy -import timm +import importlib +import unittest from typing import Dict -from utils import cosine_similarity, COSINE_THRESHOLD + +import torch +import torch_tensorrt as torchtrt +from utils import COSINE_THRESHOLD, cosine_similarity + +if importlib.util.find_spec("torchvision"): + import timm + import torchvision.models as models class TestModelToEngineToModel(unittest.TestCase): + @unittest.skipIf( + not importlib.util.find_spec("torchvision"), "torchvision not installed" + ) def test_resnet50(self): self.model = models.resnet50(pretrained=True).eval().to("cuda") self.input = torch.randn((1, 3, 224, 224)).to("cuda") @@ -37,6 +44,11 @@ def test_resnet50(self): msg=f"Resnet50 TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}", ) + @unittest.skipIf( + not importlib.util.find_spec("timm") + or not importlib.util.find_spec("torchvision"), + "timm or torchvision not installed", + ) def test_efficientnet_b0(self): self.model = ( timm.create_model("efficientnet_b0", pretrained=True).eval().to("cuda") diff --git a/tests/py/ts/model_test_case.py b/tests/py/ts/model_test_case.py index 42073ef747..bfc51df939 100644 --- a/tests/py/ts/model_test_case.py +++ b/tests/py/ts/model_test_case.py @@ -1,7 +1,7 @@ +import os import unittest + import torch -import torchvision.models as models -import os REPO_ROOT = os.path.abspath(os.getcwd()) + "/../../" diff --git a/tests/py/ts/models/test_models.py b/tests/py/ts/models/test_models.py index bae94374d8..923618a708 100644 --- a/tests/py/ts/models/test_models.py +++ b/tests/py/ts/models/test_models.py @@ -1,18 +1,24 @@ import copy +import importlib import unittest from typing import Dict -import timm 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 timm + 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 TestModels(unittest.TestCase): def test_resnet18(self): self.model = models.resnet18(pretrained=True).eval().to("cuda") @@ -39,6 +45,9 @@ def test_resnet18(self): msg=f"Resnet50 TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}", ) + @unittest.skipIf( + not importlib.util.find_spec("torchvision"), "torchvision not installed" + ) def test_mobilenet_v2(self): self.model = models.mobilenet_v2(pretrained=True).eval().to("cuda") self.input = torch.randn((1, 3, 224, 224)).to("cuda") @@ -64,6 +73,11 @@ def test_mobilenet_v2(self): msg=f"Mobilenet v2 TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}", ) + @unittest.skipIf( + not importlib.util.find_spec("timm") + or not importlib.util.find_spec("torchvision"), + "timm or torchvision not installed", + ) def test_efficientnet_b0(self): self.model = ( timm.create_model("efficientnet_b0", pretrained=True).eval().to("cuda") @@ -91,6 +105,9 @@ def test_efficientnet_b0(self): msg=f"EfficientNet-B0 TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}", ) + @unittest.skipIf( + not importlib.util.find_spec("torchvision"), "torchvision not installed" + ) def test_resnet18_half(self): self.model = models.resnet18(pretrained=True).eval().to("cuda") self.input = torch.randn((1, 3, 224, 224)).to("cuda") diff --git a/tests/py/ts/models/test_multiple_registered_engines.py b/tests/py/ts/models/test_multiple_registered_engines.py index 3dbd724755..05b6905882 100644 --- a/tests/py/ts/models/test_multiple_registered_engines.py +++ b/tests/py/ts/models/test_multiple_registered_engines.py @@ -1,18 +1,23 @@ import copy +import importlib import unittest from typing import Dict -import timm 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, "TorchScript Frontend is not available", ) +@unittest.skipIf( + not importlib.util.find_spec("torchvision"), "torchvision not installed" +) class TestModelToEngineToModel(unittest.TestCase): def test_multiple_engines(self): self.resnet18 = models.resnet18(pretrained=True).eval().to("cuda")