xref: /aosp_15_r20/external/executorch/backends/arm/test/models/test_mobilenet_v2_arm.py (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a) 
1# Copyright (c) Meta Platforms, Inc. and affiliates.
2# Copyright 2024 Arm Limited and/or its affiliates.
3# All rights reserved.
4#
5# This source code is licensed under the BSD-style license found in the
6# LICENSE file in the root directory of this source tree.
7
8import logging
9import unittest
10
11import torch
12from executorch.backends.arm.test import common
13
14from executorch.backends.arm.test.tester.arm_tester import ArmTester
15from executorch.exir import EdgeCompileConfig
16from torchvision import models, transforms
17from torchvision.models.mobilenetv2 import MobileNet_V2_Weights
18
19
20logger = logging.getLogger(__name__)
21logger.setLevel(logging.INFO)
22
23
24class TestMobileNetV2(unittest.TestCase):
25    """Tests MobileNetV2."""
26
27    mv2 = models.mobilenetv2.mobilenet_v2(weights=MobileNet_V2_Weights.DEFAULT)
28    mv2 = mv2.eval()
29    normalize = transforms.Normalize(
30        mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
31    )
32    model_inputs = (normalize(torch.randn((1, 3, 224, 224))),)
33
34    all_operators = {
35        "executorch_exir_dialects_edge__ops_aten__native_batch_norm_legit_no_training_default",
36        "executorch_exir_dialects_edge__ops_aten_add_Tensor",
37        "executorch_exir_dialects_edge__ops_aten_permute_copy_default",
38        "executorch_exir_dialects_edge__ops_aten_addmm_default",
39        "executorch_exir_dialects_edge__ops_aten_mean_dim",
40        "executorch_exir_dialects_edge__ops_aten_hardtanh_default",
41        "executorch_exir_dialects_edge__ops_aten_convolution_default",
42    }
43
44    operators_after_quantization = all_operators - {
45        "executorch_exir_dialects_edge__ops_aten__native_batch_norm_legit_no_training_default",
46    }
47
48    _edge_compile_config: EdgeCompileConfig = EdgeCompileConfig(
49        _skip_dim_order=True,  # TODO(T182928844): Delegate dim order op to backend.
50    )
51
52    def test_mv2_tosa_MI(self):
53        (
54            ArmTester(
55                self.mv2,
56                example_inputs=self.model_inputs,
57                compile_spec=common.get_tosa_compile_spec(
58                    "TOSA-0.80.0+MI", permute_memory_to_nhwc=True
59                ),
60            )
61            .export()
62            .to_edge_transform_and_lower(edge_compile_config=self._edge_compile_config)
63            .to_executorch()
64            .run_method_and_compare_outputs(inputs=self.model_inputs)
65        )
66
67    def test_mv2_tosa_BI(self):
68        (
69            ArmTester(
70                self.mv2,
71                example_inputs=self.model_inputs,
72                compile_spec=common.get_tosa_compile_spec(
73                    "TOSA-0.80.0+BI", permute_memory_to_nhwc=True
74                ),
75            )
76            .quantize()
77            .export()
78            .to_edge_transform_and_lower(edge_compile_config=self._edge_compile_config)
79            .to_executorch()
80            # atol=1.0 is a defensive upper limit
81            # TODO MLETROCH-72
82            # TODO MLETROCH-149
83            .run_method_and_compare_outputs(atol=1.0, qtol=1, inputs=self.model_inputs)
84        )
85
86    def test_mv2_u55_BI(self):
87        tester = (
88            ArmTester(
89                self.mv2,
90                example_inputs=self.model_inputs,
91                compile_spec=common.get_u55_compile_spec(permute_memory_to_nhwc=True),
92            )
93            .quantize()
94            .export()
95            .to_edge_transform_and_lower(edge_compile_config=self._edge_compile_config)
96            .to_executorch()
97            .serialize()
98        )
99        if common.is_option_enabled("corstone300"):
100            tester.run_method_and_compare_outputs(
101                atol=1.0, qtol=1, inputs=self.model_inputs, target_board="corstone-300"
102            )
103
104    def test_mv2_u85_BI(self):
105        tester = (
106            ArmTester(
107                self.mv2,
108                example_inputs=self.model_inputs,
109                compile_spec=common.get_u85_compile_spec(permute_memory_to_nhwc=True),
110            )
111            .quantize()
112            .export()
113            .to_edge_transform_and_lower(edge_compile_config=self._edge_compile_config)
114            .to_executorch()
115            .serialize()
116        )
117        if common.is_option_enabled("corstone300"):
118            tester.run_method_and_compare_outputs(
119                atol=1.0, qtol=1, inputs=self.model_inputs, target_board="corstone-320"
120            )
121