Skip to content

Commit 5985c91

Browse files
committed
Add Arm conformer unit test
Change-Id: Ic0aefae51312c533acc761d38437f4d79408a940 Signed-off-by: Fredrik Knutsson <[email protected]>
1 parent 5190106 commit 5985c91

File tree

2 files changed

+50
-1
lines changed

2 files changed

+50
-1
lines changed
Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
# Copyright 2024 Arm Limited and/or its affiliates.
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+
7+
8+
import torch
9+
from executorch.backends.arm.test import common
10+
11+
from executorch.backends.arm.test.tester.arm_tester import ArmTester
12+
from executorch.exir import EdgeCompileConfig
13+
14+
from torchaudio.models import Conformer
15+
16+
17+
dim = 16
18+
lengths = torch.randint(1, 100, (10,), dtype=torch.int32)
19+
input_data = torch.rand(10, int(lengths.max()), dim)
20+
21+
conformer = Conformer(
22+
input_dim=dim,
23+
num_heads=4,
24+
ffn_dim=64,
25+
num_layers=2,
26+
depthwise_conv_kernel_size=31,
27+
)
28+
conformer = conformer.eval()
29+
30+
_edge_compile_config: EdgeCompileConfig = EdgeCompileConfig(
31+
_skip_dim_order=True, # TODO(T182928844): Delegate dim order op to backend.
32+
)
33+
34+
35+
def test_conformer_tosa_MI():
36+
(
37+
ArmTester(
38+
conformer,
39+
example_inputs=(input_data, lengths),
40+
compile_spec=common.get_tosa_compile_spec(tosa_version="TOSA-0.80+MI"),
41+
)
42+
.export()
43+
.dump_operator_distribution("ta_conformer_ops_export.txt")
44+
.to_edge() # todo: .to_edge_transform_and_lower(edge_compile_config=_edge_compile_config)
45+
.dump_operator_distribution("ta_conformer_ops_toedge.txt")
46+
# .to_executorch()
47+
# .run_method_and_compare_outputs(inputs=(input_data, lengths))
48+
# .serialize(PATH)
49+
)

third-party/ao

Submodule ao updated 378 files

0 commit comments

Comments
 (0)