Arm backend: Add FP16 VGF coverage for DeiT Tiny (pytorch#21716)

Add support for FP16 to DeiT Tiny. Fixture is added to download the FP32
model once and use it for both FP32 inference and also quantized (FP16)
so that two runs can be compared later for various purposes.

Signed-off-by: Sangwon Ha <sangwon.ha@arm.com>

Author: FabulousSuperDude

backends/arm/test/models/test_deit_tiny_arm.py

@@ -5,6 +5,7 @@
 
 from typing import Tuple
 
+import pytest
 import timm  # type: ignore[import-untyped]
 
 import torch
@@ -25,30 +26,66 @@
 )
 from torchvision import transforms  # type: ignore[import-untyped]
 
-deit_tiny = timm.models.deit.deit_tiny_patch16_224(pretrained=True)
-
-deit_tiny.eval()
-
 normalize = transforms.Normalize(
     mean=IMAGENET_INCEPTION_MEAN, std=IMAGENET_INCEPTION_STD
 )
 model_inputs = (normalize(torch.rand((1, 3, 224, 224))),)
+image_fp16 = torch.randn((1, 3, 224, 224), generator=torch.Generator().manual_seed(0))
+model_inputs_fp16 = (normalize(image_fp16).to(torch.float16),)
 
 input_t = Tuple[torch.Tensor]
 
 
-def test_deit_tiny_tosa_FP():
-    pipeline = TosaPipelineFP[input_t](
-        deit_tiny,
+@pytest.fixture(scope="module")
+def deit_tiny():
+    return timm.models.deit.deit_tiny_patch16_224(pretrained=True).eval()
+
+
+@pytest.fixture(scope="module")
+def get_deit_model(deit_tiny):
+    models = {torch.float32: deit_tiny}
+
+    def get_model(dtype):
+        if dtype not in models:
+            model = timm.models.deit.deit_tiny_patch16_224(pretrained=False)
+            model.load_state_dict(deit_tiny.state_dict())
+            models[dtype] = model.to(dtype).eval()
+        return models[dtype]
+
+    return get_model
+
+
+fp_test_data = [
+    pytest.param(
+        torch.float32,
         model_inputs,
+        {"atol": 1e-3, "rtol": 1e-3},
+        id="fp32",
+    ),
+    pytest.param(
+        torch.float16,
+        model_inputs_fp16,
+        {"atol": 8e-2, "rtol": 8e-2},
+        marks=pytest.mark.slow,
+        id="fp16",
+    ),
+]
+
+
+@pytest.mark.parametrize("dtype,inputs,pipeline_kwargs", fp_test_data)
+def test_deit_tiny_tosa_FP(dtype, inputs, pipeline_kwargs, get_deit_model):
+    pipeline = TosaPipelineFP[input_t](
+        get_deit_model(dtype),
+        inputs,
         aten_op=[],
         exir_op=[],
         use_to_edge_transform_and_lower=True,
+        **pipeline_kwargs,
     )
     pipeline.run()
 
 
-def test_deit_tiny_tosa_INT():
+def test_deit_tiny_tosa_INT(deit_tiny):
     pipeline = TosaPipelineINT[input_t](
         deit_tiny,
         model_inputs,
@@ -63,7 +100,7 @@ def test_deit_tiny_tosa_INT():
     pipeline.run()
 
 
-def test_deit_tiny_u55_INT():
+def test_deit_tiny_u55_INT(deit_tiny):
     pipeline = EthosU55PipelineINT[input_t](
         deit_tiny,
         model_inputs,
@@ -81,7 +118,7 @@ def test_deit_tiny_u55_INT():
 
 
 @common.XfailIfNoCorstone320
-def test_deit_tiny_u85_INT():
+def test_deit_tiny_u85_INT(deit_tiny):
     pipeline = EthosU85PipelineINT[input_t](
         deit_tiny,
         model_inputs,
@@ -95,7 +132,7 @@ def test_deit_tiny_u85_INT():
 
 
 @common.SkipIfNoModelConverter
-def test_deit_tiny_vgf_quant():
+def test_deit_tiny_vgf_quant(deit_tiny):
     pipeline = VgfPipeline[input_t](
         deit_tiny,
         model_inputs,
@@ -110,13 +147,15 @@ def test_deit_tiny_vgf_quant():
 
 
 @common.SkipIfNoModelConverter
-def test_deit_tiny_vgf_no_quant():
+@pytest.mark.parametrize("dtype,inputs,pipeline_kwargs", fp_test_data)
+def test_deit_tiny_vgf_no_quant(dtype, inputs, pipeline_kwargs, get_deit_model):
     pipeline = VgfPipeline[input_t](
-        deit_tiny,
-        model_inputs,
+        get_deit_model(dtype),
+        inputs,
         aten_op=[],
         exir_op=[],
         use_to_edge_transform_and_lower=True,
         quantize=False,
+        **pipeline_kwargs,
     )
     pipeline.run()