xref: /aosp_15_r20/external/executorch/examples/qualcomm/oss_scripts/squeezenet.py (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a)
1# Copyright (c) Qualcomm Innovation Center, Inc.
2# All rights reserved
3#
4# This source code is licensed under the BSD-style license found in the
5# LICENSE file in the root directory of this source tree.
6
7import json
8import os
9from multiprocessing.connection import Client
10
11import numpy as np
12
13import torch
14from executorch.backends.qualcomm.quantizer.quantizer import QuantDtype
15from executorch.examples.qualcomm.utils import (
16    build_executorch_binary,
17    get_imagenet_dataset,
18    make_output_dir,
19    parse_skip_delegation_node,
20    setup_common_args_and_variables,
21    SimpleADB,
22    topk_accuracy,
23)
24
25
26def main(args):
27    skip_node_id_set, skip_node_op_set = parse_skip_delegation_node(args)
28
29    # ensure the working directory exist.
30    os.makedirs(args.artifact, exist_ok=True)
31
32    if not args.compile_only and args.device is None:
33        raise RuntimeError(
34            "device serial is required if not compile only. "
35            "Please specify a device serial by -s/--device argument."
36        )
37
38    data_num = 100
39    inputs, targets, input_list = get_imagenet_dataset(
40        dataset_path=f"{args.dataset}",
41        data_size=data_num,
42        image_shape=(256, 256),
43        crop_size=224,
44    )
45    pte_filename = "squeezenet_qnn"
46    instance = torch.hub.load(
47        "pytorch/vision:v0.13.0",
48        "squeezenet1_1",
49        weights="SqueezeNet1_1_Weights.DEFAULT",
50    )
51    build_executorch_binary(
52        instance.eval(),
53        (torch.randn(1, 3, 224, 224),),
54        args.model,
55        f"{args.artifact}/{pte_filename}",
56        inputs,
57        skip_node_id_set=skip_node_id_set,
58        skip_node_op_set=skip_node_op_set,
59        quant_dtype=QuantDtype.use_8a8w,
60        shared_buffer=args.shared_buffer,
61    )
62
63    if args.compile_only:
64        return
65
66    adb = SimpleADB(
67        qnn_sdk=os.getenv("QNN_SDK_ROOT"),
68        build_path=f"{args.build_folder}",
69        pte_path=f"{args.artifact}/{pte_filename}.pte",
70        workspace=f"/data/local/tmp/executorch/{pte_filename}",
71        device_id=args.device,
72        host_id=args.host,
73        soc_model=args.model,
74    )
75    adb.push(inputs=inputs, input_list=input_list)
76    adb.execute()
77
78    # collect output data
79    output_data_folder = f"{args.artifact}/outputs"
80    make_output_dir(output_data_folder)
81
82    adb.pull(output_path=args.artifact)
83
84    # top-k analysis
85    predictions = []
86    for i in range(data_num):
87        predictions.append(
88            np.fromfile(
89                os.path.join(output_data_folder, f"output_{i}_0.raw"), dtype=np.float32
90            )
91        )
92
93    k_val = [1, 5]
94    topk = [topk_accuracy(predictions, targets, k).item() for k in k_val]
95    if args.ip and args.port != -1:
96        with Client((args.ip, args.port)) as conn:
97            conn.send(json.dumps({f"top_{k}": topk[i] for i, k in enumerate(k_val)}))
98    else:
99        for i, k in enumerate(k_val):
100            print(f"top_{k}->{topk[i]}%")
101
102
103if __name__ == "__main__":
104    parser = setup_common_args_and_variables()
105
106    parser.add_argument(
107        "-d",
108        "--dataset",
109        help=(
110            "path to the validation folder of ImageNet dataset. "
111            "e.g. --dataset imagenet-mini/val "
112            "for https://www.kaggle.com/datasets/ifigotin/imagenetmini-1000)"
113        ),
114        type=str,
115        required=True,
116    )
117
118    parser.add_argument(
119        "-a",
120        "--artifact",
121        help="path for storing generated artifacts by this example. "
122        "Default ./squeezenet",
123        default="./squeezenet",
124        type=str,
125    )
126
127    args = parser.parse_args()
128    try:
129        main(args)
130    except Exception as e:
131        if args.ip and args.port != -1:
132            with Client((args.ip, args.port)) as conn:
133                conn.send(json.dumps({"Error": str(e)}))
134        else:
135            raise Exception(e)
136