add qnn tests

Author: cccclai

tests/models/test_modeling_bert.py

@@ -22,12 +22,12 @@
 import pytest
 import torchao
 from executorch.extension.pybindings.portable_lib import ExecuTorchModule
+
+from optimum.executorch import ExecuTorchModelForMaskedLM
 from packaging.version import parse
 from transformers import AutoTokenizer
 from transformers.testing_utils import slow
 
-from optimum.executorch import ExecuTorchModelForMaskedLM
-
 
 @pytest.mark.skipif(
     parse(torchao.__version__) < parse("0.11.0.dev0"),
@@ -70,7 +70,9 @@ def _helper_bert_fill_mask(self, recipe: str):
         tokenizer = AutoTokenizer.from_pretrained(model_id)
 
         # Test fetching and lowering the model to ExecuTorch
-        model = ExecuTorchModelForMaskedLM.from_pretrained(model_id=model_id, recipe=recipe)
+        model = ExecuTorchModelForMaskedLM.from_pretrained(
+            model_id=model_id, recipe=recipe
+        )
         self.assertIsInstance(model, ExecuTorchModelForMaskedLM)
         self.assertIsInstance(model.model, ExecuTorchModule)
 
@@ -85,9 +87,14 @@ def _helper_bert_fill_mask(self, recipe: str):
         # Test inference using ExecuTorch model
         exported_outputs = model.forward(inputs["input_ids"], inputs["attention_mask"])
         predicted_masks = tokenizer.decode(exported_outputs[0, 4].topk(5).indices)
-        logging.info(f"\nInput text:\n\t{input_text}\nPredicted masks:\n\t{predicted_masks}")
+        logging.info(
+            f"\nInput text:\n\t{input_text}\nPredicted masks:\n\t{predicted_masks}"
+        )
         self.assertTrue(
-            any(word in predicted_masks for word in ["capital", "center", "heart", "birthplace"]),
+            any(
+                word in predicted_masks
+                for word in ["capital", "center", "heart", "birthplace"]
+            ),
             f"Exported model predictions {predicted_masks} don't contain any of the most common expected words",
         )
 
@@ -101,3 +108,7 @@ def test_bert_fill_mask(self):
     @pytest.mark.portable
     def test_bert_fill_mask_portable(self):
         self._helper_bert_fill_mask("portable")
+
+    @pytest.mark.run_slow
+    def test_bert_fill_mask_qnn(self):
+        self._helper_bert_fill_mask(recipe="qnn_fp16_SM8650")

tests/models/test_modeling_cvt.py

@@ -21,10 +21,10 @@
 import pytest
 import torch
 from executorch.extension.pybindings.portable_lib import ExecuTorchModule
-from transformers import AutoConfig, AutoModelForImageClassification
-from transformers.testing_utils import slow
 
 from optimum.executorch import ExecuTorchModelForImageClassification
+from transformers import AutoConfig, AutoModelForImageClassification
+from transformers.testing_utils import slow
 
 from ..utils import check_close_recursively
 
@@ -58,11 +58,15 @@ def _helper_cvt_image_classification(self, recipe: str):
         pixel_values = torch.rand(batch_size, num_channels, height, width)
 
         # Test fetching and lowering the model to ExecuTorch
-        et_model = ExecuTorchModelForImageClassification.from_pretrained(model_id=model_id, recipe=recipe)
+        et_model = ExecuTorchModelForImageClassification.from_pretrained(
+            model_id=model_id, recipe=recipe
+        )
         self.assertIsInstance(et_model, ExecuTorchModelForImageClassification)
         self.assertIsInstance(et_model.model, ExecuTorchModule)
 
-        eager_model = AutoModelForImageClassification.from_pretrained(model_id).eval().to("cpu")
+        eager_model = (
+            AutoModelForImageClassification.from_pretrained(model_id).eval().to("cpu")
+        )
         with torch.no_grad():
             eager_output = eager_model(pixel_values)
             et_output = et_model.forward(pixel_values)
@@ -80,3 +84,8 @@ def test_cvt_image_classification(self):
     @pytest.mark.portable
     def test_cvt_image_classification_portable(self):
         self._helper_cvt_image_classification(recipe="portable")
+
+    @slow
+    @pytest.mark.run_slow
+    def test_cvt_image_classification_qnn(self):
+        self._helper_cvt_image_classification(recipe="qnn_fp16_SM8650")

tests/models/test_modeling_deit.py

@@ -21,10 +21,10 @@
 import pytest
 import torch
 from executorch.extension.pybindings.portable_lib import ExecuTorchModule
-from transformers import AutoConfig, AutoModelForImageClassification
-from transformers.testing_utils import slow
 
 from optimum.executorch import ExecuTorchModelForImageClassification
+from transformers import AutoConfig, AutoModelForImageClassification
+from transformers.testing_utils import slow
 
 from ..utils import check_close_recursively
 
@@ -58,11 +58,15 @@ def _helper_deit_image_classification(self, recipe: str):
         pixel_values = torch.rand(batch_size, num_channels, height, width)
 
         # Test fetching and lowering the model to ExecuTorch
-        et_model = ExecuTorchModelForImageClassification.from_pretrained(model_id=model_id, recipe=recipe)
+        et_model = ExecuTorchModelForImageClassification.from_pretrained(
+            model_id=model_id, recipe=recipe
+        )
         self.assertIsInstance(et_model, ExecuTorchModelForImageClassification)
         self.assertIsInstance(et_model.model, ExecuTorchModule)
 
-        eager_model = AutoModelForImageClassification.from_pretrained(model_id).eval().to("cpu")
+        eager_model = (
+            AutoModelForImageClassification.from_pretrained(model_id).eval().to("cpu")
+        )
         with torch.no_grad():
             eager_output = eager_model(pixel_values)
             et_output = et_model.forward(pixel_values)
@@ -80,3 +84,8 @@ def test_deit_image_classification(self):
     @pytest.mark.portable
     def test_deit_image_classification_portable(self):
         self._helper_deit_image_classification(recipe="portable")
+
+    @slow
+    @pytest.mark.run_slow
+    def test_deit_image_classification_qnn(self):
+        self._helper_deit_image_classification(recipe="qnn_fp16_SM8650")