Skip to content

Commit 957915f

Browse files
Arm Backend: Add visualization script for Arm models (#14257)
This PR enables the visualisation of a tosa flatbuffer file in the arm `run.sh` script. The `model_explorer` command line flag is added which launches model explorer, with the tosa adapter plugin - https://github.com/Arm/tosa-adapter-model-explorer. cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 Co-authored-by: Robert Taylor <[email protected]>
1 parent dfd7f2a commit 957915f

File tree

3 files changed

+42
-0
lines changed

3 files changed

+42
-0
lines changed

backends/arm/requirements-arm-tosa.txt

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,5 +5,7 @@
55

66
ml_dtypes == 0.5.1
77
flatbuffers == 24.3.25
8+
tosa-adapter-model-explorer == 0.0.1
9+
ai-edge-model-explorer >= 0.1.16
810

911
tosa-tools @ git+https://git.gitlab.arm.com/tosa/[email protected]

examples/arm/run.sh

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ scratch_dir_set=false
4141
toolchain=arm-none-eabi-gcc
4242
select_ops_list="aten::_softmax.out"
4343
qdq_fusion_op=false
44+
model_explorer=false
4445

4546
function help() {
4647
echo "Usage: $(basename $0) [options]"
@@ -71,6 +72,7 @@ function help() {
7172
echo " --et_build_root=<FOLDER> Executorch build output root folder to use, defaults to ${et_build_root}"
7273
echo " --scratch-dir=<FOLDER> Path to your Ethos-U scrach dir if you not using default ${ethos_u_scratch_dir}"
7374
echo " --qdq_fusion_op Enable QDQ fusion op"
75+
echo " --model_explorer Generate and open a visual graph of the compiled model."
7476
exit 0
7577
}
7678

@@ -99,6 +101,7 @@ for arg in "$@"; do
99101
--et_build_root=*) et_build_root="${arg#*=}";;
100102
--scratch-dir=*) ethos_u_scratch_dir="${arg#*=}" ; scratch_dir_set=true ;;
101103
--qdq_fusion_op) qdq_fusion_op=true;;
104+
--model_explorer) model_explorer=true ;;
102105
*)
103106
;;
104107
esac
@@ -326,6 +329,11 @@ for i in "${!test_model[@]}"; do
326329
fi
327330
set +x
328331
fi
332+
333+
if [ "$model_explorer" = true ]; then
334+
tosa_flatbuffer_path=$(find ${output_folder} -name "*TOSA*.tosa" | head -n 1)
335+
python3 ${script_dir}/visualize.py ${tosa_flatbuffer_path}
336+
fi
329337
done
330338

331339
exit 0

examples/arm/visualize.py

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
# Copyright 2025 Arm Limited and/or its affiliates.
2+
#
3+
# This source code is licensed under the BSD-style license found in the
4+
# LICENSE file in the root directory of this source tree.
5+
6+
import argparse
7+
8+
import model_explorer
9+
10+
from executorch.devtools.visualization.visualization_utils import (
11+
visualize_model_explorer,
12+
)
13+
14+
15+
def main() -> None:
16+
parser = argparse.ArgumentParser(
17+
description="Visualize a model using model explorer."
18+
)
19+
parser.add_argument("model_path", type=str, help="Path to the model file.")
20+
args = parser.parse_args()
21+
22+
config = model_explorer.config()
23+
(config.add_model_from_path(args.model_path))
24+
25+
visualize_model_explorer(
26+
config=config,
27+
extensions=["tosa_adapter_model_explorer"],
28+
)
29+
30+
31+
if __name__ == "__main__":
32+
main()

0 commit comments

Comments
 (0)