build: Guarded torchvision calls with available_packages

Signed-off-by: Brandon Groth <[email protected]>

Author: BrandonGroth

fms_mo/fx/dynamo_utils.py

@@ -29,6 +29,7 @@
     get_target_op_from_mod_or_str,
     get_target_op_from_node,
 )
+from fms_mo.utils.import_utils import available_packages
 
 logger = logging.getLogger(__name__)
 
@@ -1133,7 +1134,6 @@ def cus_backend_model_analyzer(
     from functools import partial
 
     # Third Party
-    from torchvision.models import VisionTransformer
     from transformers import PreTrainedModel
 
     if issubclass(type(model), torch.nn.Module):
@@ -1145,7 +1145,16 @@ def cus_backend_model_analyzer(
         model_to_be_traced = model
         model_param_size = 999
 
-    is_transformers = issubclass(type(model), (PreTrainedModel, VisionTransformer))
+    transformer_model_classes = (PreTrainedModel,)
+
+    if available_packages["torchvision"]:
+        # Third Party
+        # pylint: disable = import-error
+        from torchvision.models import VisionTransformer
+
+        transformer_model_classes += (VisionTransformer,)
+
+    is_transformers = issubclass(type(model), transformer_model_classes)
     if model_param_size > 1:
         # Standard
         import sys
@@ -1188,11 +1197,13 @@ def call_seq_hook(mod, *_args, **_kwargs):
 
         # only add last layer
         qcfg["qskip_layer_name"] += [qcfg["mod_call_seq"][-1]]
-        # unless it's a ViT, skip first Conv as well
-        if issubclass(type(model), VisionTransformer) and isinstance(
-            model.get_submodule(qcfg["mod_call_seq"][0]), torch.nn.Conv2d
-        ):
-            qcfg["qskip_layer_name"] += [qcfg["mod_call_seq"][0]]
+
+        if available_packages["torchvision"]:
+            # unless it's a ViT, skip first Conv as well
+            if issubclass(type(model), VisionTransformer) and isinstance(
+                model.get_submodule(qcfg["mod_call_seq"][0]), torch.nn.Conv2d
+            ):
+                qcfg["qskip_layer_name"] += [qcfg["mod_call_seq"][0]]
 
     with torch.no_grad():
         model_opt = torch.compile(
@@ -1271,21 +1282,23 @@ def qbmm_auto_check(_mod, *_args, **_kwargs):
     # c) identify RPN/FPN
     # TODO this hack only works for torchvision models. will use find_rpn_fpn_gm()
 
-    # Third Party
-    from torchvision.models.detection.rpn import RegionProposalNetwork
-    from torchvision.ops import FeaturePyramidNetwork
-
-    rpnfpn_prefix = []
-    rpnfpn_convs = []
-    for n, m in model.named_modules():
-        if isinstance(m, (FeaturePyramidNetwork, RegionProposalNetwork)):
-            rpnfpn_prefix.append(n)
-        if isinstance(m, torch.nn.Conv2d) and any(
-            n.startswith(p) for p in rpnfpn_prefix
-        ):
-            rpnfpn_convs.append(n)
-            if n not in qcfg["qskip_layer_name"]:
-                qcfg["qskip_layer_name"].append(n)
+    if available_packages["torchvision"]:
+        # Third Party
+        # pylint: disable = import-error
+        from torchvision.models.detection.rpn import RegionProposalNetwork
+        from torchvision.ops import FeaturePyramidNetwork
+
+        rpnfpn_prefix = []
+        rpnfpn_convs = []
+        for n, m in model.named_modules():
+            if isinstance(m, (FeaturePyramidNetwork, RegionProposalNetwork)):
+                rpnfpn_prefix.append(n)
+            if isinstance(m, torch.nn.Conv2d) and any(
+                n.startswith(p) for p in rpnfpn_prefix
+            ):
+                rpnfpn_convs.append(n)
+                if n not in qcfg["qskip_layer_name"]:
+                    qcfg["qskip_layer_name"].append(n)
 
     if qcfg["N_backend_called"] > 1:
         logger.warning(

tests/models/conftest.py

@@ -23,8 +23,6 @@
 
 # Third Party
 from torch.utils.data import DataLoader, TensorDataset
-from torchvision.io import read_image
-from torchvision.models import ResNet50_Weights, ViT_B_16_Weights, resnet50, vit_b_16
 from transformers import (
     BertConfig,
     BertModel,
@@ -43,6 +41,7 @@
 # fms_mo imports
 from fms_mo import qconfig_init
 from fms_mo.modules import QLSTM, QBmm, QConv2d, QConvTranspose2d, QLinear
+from fms_mo.utils.import_utils import available_packages
 from fms_mo.utils.qconfig_utils import get_mx_specs_defaults, set_mx_specs
 
 ########################
@@ -1123,75 +1122,82 @@ def required_pair(request):
 # Vision Model Fixtures #
 #########################
 
-# Create img
-# downloaded from torchvision github (vision/test/assets/encoder_jpeg/ directory)
-img = read_image(
-    os.path.realpath(
-        os.path.join(os.path.dirname(__file__), "grace_hopper_517x606.jpg")
-    )
-)
-
 
-# Create resnet/vit batch fixtures from weights
-def prepocess_img(image, weights):
-    """
-    Preprocess an image w/ a weights.transform()
-
-    Args:
-        img (torch.FloatTensor): Image data
-        weights (torchvision.models): Weight object
-
-    Returns:
-        torch.FloatTensor: Preprocessed image
-    """
-    preprocess = weights.transforms()
-    batch = preprocess(image).unsqueeze(0)
-    return batch
+if available_packages["torchvision"]:
+    # Third Party
+    # pylint: disable = import-error
+    from torchvision.io import read_image
+    from torchvision.models import (
+        ResNet50_Weights,
+        ViT_B_16_Weights,
+        resnet50,
+        vit_b_16,
+    )
 
+    # Create img
+    # downloaded from torchvision github (vision/test/assets/encoder_jpeg/ directory)
+    img_tv = read_image(
+        os.path.realpath(
+            os.path.join(os.path.dirname(__file__), "grace_hopper_517x606.jpg")
+        )
+    )
 
-@pytest.fixture(scope="session")
-def batch_resnet():
-    """
-    Preprocess an image w/ Resnet weights.transform()
+    # Create resnet/vitbatch fixtures from weights
+    def prepocess_img(image, weights):
+        """
+        Preprocess an image w/ a weights.transform()
 
-    Returns:
-        torch.FloatTensor: Preprocessed image
-    """
-    return prepocess_img(img, ResNet50_Weights.IMAGENET1K_V2)
+        Args:
+            img_tv (torch.FloatTensor): Image data
+            weights (torchvision.models): Weight object
 
+        Returns:
+            torch.FloatTensor: Preprocessed image
+        """
+        preprocess = weights.transforms()
+        batch = preprocess(image).unsqueeze(0)
+        return batch
 
-@pytest.fixture(scope="session")
-def batch_vit():
-    """
-    Preprocess an image w/ ViT weights.transform()
+    @pytest.fixture(scope="session")
+    def batch_resnet():
+        """
+        Preprocess an image w/ Resnet weights.transform()
 
-    Returns:
-        torch.FloatTensor: Preprocessed image
-    """
-    return prepocess_img(img, ViT_B_16_Weights.IMAGENET1K_V1)
+        Returns:
+            torch.FloatTensor: Preprocessed image
+        """
+        return prepocess_img(img_tv, ResNet50_Weights.IMAGENET1K_V2)
 
+    @pytest.fixture(scope="session")
+    def batch_vit():
+        """
+        Preprocess an image w/ ViT weights.transform()
 
-# Create resnet/vit model fixtures from weights
-@pytest.fixture(scope="function")
-def model_resnet():
-    """
-    Create Resnet50 model + weights
+        Returns:
+            torch.FloatTensor: Preprocessed image
+        """
+        return prepocess_img(img_tv, ViT_B_16_Weights.IMAGENET1K_V1)
 
-    Returns:
-        torchvision.models.resnet.ResNet: Resnet50 model
-    """
-    return resnet50(weights=ResNet50_Weights.IMAGENET1K_V2)
+    # Create resnet/vit model fixtures from weights
+    @pytest.fixture(scope="function")
+    def model_resnet():
+        """
+        Create Resnet50 model + weights
 
+        Returns:
+            torchvision.models.resnet.ResNet: Resnet50 model
+        """
+        return resnet50(weights=ResNet50_Weights.IMAGENET1K_V2)
 
-@pytest.fixture(scope="function")
-def model_vit():
-    """
-    Create ViT model + weights
+    @pytest.fixture(scope="function")
+    def model_vit():
+        """
+        Create ViT model + weights
 
-    Returns:
-        torchvision.models.vision_transformer.VisionTransformer: ViT model
-    """
-    return vit_b_16(weights=ViT_B_16_Weights.IMAGENET1K_V1)
+        Returns:
+            torchvision.models.vision_transformer.VisionTransformer: ViT model
+        """
+        return vit_b_16(weights=ViT_B_16_Weights.IMAGENET1K_V1)
 
 
 #######################

tests/models/test_qmodelprep.py

@@ -19,13 +19,13 @@
 # Third Party
 import pytest
 import torch
-import torchvision
 import transformers
 
 # Local
 # fms_mo imports
 from fms_mo import qconfig_init, qmodel_prep
 from fms_mo.prep import has_quantized_module
+from fms_mo.utils.import_utils import available_packages
 from fms_mo.utils.utils import patch_torch_bmm
 from tests.models.test_model_utils import count_qmodules, delete_file, qmodule_error
 
@@ -159,8 +159,12 @@ def test_config_fp32_qmodes(
 ###########################
 
 
+@pytest.mark.skipif(
+    not available_packages["torchvision"],
+    reason="Requires torchvision",
+)
 def test_resnet50_torchscript(
-    model_resnet: torchvision.models.resnet.ResNet,
+    model_resnet,
     batch_resnet: torch.FloatTensor,
     config_int8: dict,
 ):
@@ -177,8 +181,12 @@ def test_resnet50_torchscript(
     qmodule_error(model_resnet, 6, 48)
 
 
+@pytest.mark.skipif(
+    not available_packages["torchvision"],
+    reason="Requires torchvision",
+)
 def test_resnet50_dynamo(
-    model_resnet: torchvision.models.resnet.ResNet,
+    model_resnet,
     batch_resnet: torch.FloatTensor,
     config_int8: dict,
 ):
@@ -195,8 +203,12 @@ def test_resnet50_dynamo(
     qmodule_error(model_resnet, 6, 48)
 
 
+@pytest.mark.skipif(
+    not available_packages["torchvision"],
+    reason="Requires torchvision",
+)
 def test_resnet50_dynamo_layers(
-    model_resnet: torchvision.models.resnet.ResNet,
+    model_resnet,
     batch_resnet: torch.FloatTensor,
     config_int8: dict,
 ):
@@ -216,8 +228,12 @@ def test_resnet50_dynamo_layers(
 
 
 # Vision Transformer tests
+@pytest.mark.skipif(
+    not available_packages["torchvision"],
+    reason="Requires torchvision",
+)
 def test_vit_torchscript(
-    model_vit: torchvision.models.vision_transformer.VisionTransformer,
+    model_vit,
     batch_vit: torch.FloatTensor,
     config_int8: dict,
 ):
@@ -234,8 +250,12 @@ def test_vit_torchscript(
     qmodule_error(model_vit, 2, 36)
 
 
+@pytest.mark.skipif(
+    not available_packages["torchvision"],
+    reason="Requires torchvision",
+)
 def test_vit_dynamo(
-    model_vit: torchvision.models.vision_transformer.VisionTransformer,
+    model_vit,
     batch_vit: torch.FloatTensor,
     config_int8: dict,
 ):