mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-21 19:18:55 +00:00
97cc40d75a
### Description Add conformer-transducer model type to optimizer. This PR adds pattern matches for attention shown below: Unfused attention:  Fused attention: 
69 lines
2.6 KiB
Python
69 lines
2.6 KiB
Python
# -------------------------------------------------------------------------
|
|
# Copyright (c) Microsoft Corporation. All rights reserved.
|
|
# Licensed under the MIT License. See License.txt in the project root for
|
|
# license information.
|
|
# --------------------------------------------------------------------------
|
|
|
|
import os
|
|
import unittest
|
|
|
|
import onnx
|
|
from conformer_model_generator import create_conformer_attention
|
|
from parity_utilities import find_transformers_source
|
|
|
|
if find_transformers_source():
|
|
from fusion_options import FusionOptions
|
|
from onnx_model import OnnxModel
|
|
from optimizer import optimize_model
|
|
else:
|
|
from onnxruntime.transformers.fusion_options import FusionOptions
|
|
from onnxruntime.transformers.onnx_model import OnnxModel
|
|
from onnxruntime.transformers.optimizer import optimize_model
|
|
|
|
|
|
class TestFusion(unittest.TestCase):
|
|
def verify_fusion(self, optimized_model, expected_model_filename):
|
|
optimized_model.topological_sort(is_deterministic=True)
|
|
|
|
expected_model_path = os.path.join(
|
|
os.path.dirname(__file__), "test_data", "models", "conformer", expected_model_filename
|
|
)
|
|
print("Expected model path = ", expected_model_path)
|
|
expected_model = OnnxModel(onnx.load(expected_model_path))
|
|
expected_model.topological_sort(is_deterministic=True)
|
|
|
|
nodes = optimized_model.model.graph.node
|
|
self.assertEqual(len(nodes), len(expected_model.model.graph.node))
|
|
|
|
for i in range(len(nodes)):
|
|
self.assertEqual(nodes[i], expected_model.model.graph.node[i])
|
|
|
|
for expected_initializer in expected_model.model.graph.initializer:
|
|
print("Expected initializer initial = ", expected_initializer.name)
|
|
self.assertTrue(
|
|
OnnxModel.has_same_value(
|
|
optimized_model.get_initializer(expected_initializer.name), expected_initializer
|
|
)
|
|
)
|
|
|
|
def test_ct_mha_fusion(self):
|
|
num_heads = 8
|
|
hidden_size = 512
|
|
model = create_conformer_attention(num_heads=num_heads, hidden_size=hidden_size, add_before_layernorm=False)
|
|
dir = "."
|
|
model_path = os.path.join(dir, "conformer_self_mha.onnx")
|
|
onnx.save(model, model_path)
|
|
options = FusionOptions("conformer")
|
|
optimized_model = optimize_model(
|
|
model_path,
|
|
model_type="conformer",
|
|
num_heads=num_heads,
|
|
hidden_size=hidden_size,
|
|
optimization_options=options,
|
|
)
|
|
os.remove(model_path)
|
|
self.verify_fusion(optimized_model, "conformer_self_mha_fused.onnx")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|