test: Added non-torchvision resnet + vit tests

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

Author: BrandonGroth

tests/models/conftest.py

@@ -22,8 +22,11 @@
 import os
 
 # Third Party
+from PIL import Image  # pylint: disable=import-error
 from torch.utils.data import DataLoader, TensorDataset
 from transformers import (
+    AutoImageProcessor,
+    AutoModelForImageClassification,
     BertConfig,
     BertModel,
     BertTokenizer,
@@ -1200,6 +1203,81 @@ def model_vit():
         return vit_b_16(weights=ViT_B_16_Weights.IMAGENET1K_V1)
 
 
+img = Image.open(
+    os.path.realpath(
+        os.path.join(os.path.dirname(__file__), "grace_hopper_517x606.jpg")
+    )
+).convert("RGB")
+
+
+def process_img(
+    pretrained_model: str,
+    input_img: Image.Image,
+):
+    """
+    Process an image w/ AutoImageProcessor
+
+    Args:
+        processor (AutoImageProcessor): Processor weights for pretrained model
+        pretrained_model (str): Weight object
+        input_img (Image.Image): Image data
+
+    Returns:
+        torch.FloatTensor: Processed image
+    """
+    img_processor = AutoImageProcessor.from_pretrained(pretrained_model)
+    batch_dict = img_processor(images=input_img, return_tensor="pt", use_fast=False)
+    # Data is {pixel_values: numpy_array[0]=data} w/ tensor.shape [C,W,H]
+    # Needs to be [1,C,W,H] -> unsqueeze(0)
+    return torch.from_numpy(batch_dict["pixel_values"][0]).unsqueeze(0)
+
+
+@pytest.fixture(scope="function")
+def batch_resnet18():
+    """
+    Preprocess an image w/ ms resnet18 processor
+
+    Returns:
+        torch.FloatTensor: Preprocessed image
+    """
+    return process_img("microsoft/resnet-18", img)
+
+
+@pytest.fixture(scope="function")
+def model_resnet18():
+    """
+    Create MS ResNet18 model + weights
+
+    Returns:
+        AutoModelForImageClassification: Resnet18 model
+    """
+    return AutoModelForImageClassification.from_pretrained("microsoft/resnet-18")
+
+
+@pytest.fixture(scope="function")
+def batch_vit_base():
+    """
+    Preprocess an image w/ Google ViT-base processor
+
+    Returns:
+        torch.FloatTensor: Preprocessed image
+    """
+    return process_img("google/vit-base-patch16-224", img)
+
+
+@pytest.fixture(scope="function")
+def model_vit_base():
+    """
+    Create Google ViT-base model + weights
+
+    Returns:
+        AutoModelForImageClassification: Google ViT-base model
+    """
+    return AutoModelForImageClassification.from_pretrained(
+        "google/vit-base-patch16-224"
+    )
+
+
 #######################
 # BERT Model Fixtures #
 #######################

tests/models/test_qmodelprep.py

@@ -272,6 +272,42 @@ def test_vit_dynamo(
     qmodule_error(model_vit, 2, 36)
 
 
+def test_resnet18(
+    model_resnet18,
+    batch_resnet18,
+    config_int8: dict,
+):
+    """
+    Perform int8 quantization on ResNet-18 w/ Dynamo tracer
+
+    Args:
+        model_resnet18 (AutoModelForImageClassification): Resnet18 model + weights
+        batch_resnet18 (torch.FloatTensor): Batch image data for Resnet18
+        config (dict): Recipe Config w/ int8 settings
+    """
+    # Run qmodel_prep w/ Dynamo tracer
+    qmodel_prep(model_resnet18, batch_resnet18, config_int8, use_dynamo=True)
+    qmodule_error(model_resnet18, 4, 17)
+
+
+def test_vit_base(
+    model_vit_base,
+    batch_vit_base,
+    config_int8: dict,
+):
+    """
+    Perform int8 quantization on ViT-base w/ Dynamo tracer
+
+    Args:
+        model_vit_base (AutoModelForImageClassification): Resnet18 model + weights
+        batch_vit_base (torch.FloatTensor): Batch image data for Resnet18
+        config (dict): Recipe Config w/ int8 settings
+    """
+    # Run qmodel_prep w/ Dynamo tracer
+    qmodel_prep(model_vit_base, batch_vit_base, config_int8, use_dynamo=True)
+    qmodule_error(model_vit_base, 1, 73)
+
+
 def test_bert_dynamo(
     model_bert: transformers.models.bert.modeling_bert.BertModel,
     input_bert: torch.FloatTensor,