Support symbolic shape infer in transformers tool (#6899)

* fusion support runtime edge shape checking

* trim ctor

* add test

* fix

* Update test_shape_infer_helper.py

* use torch input size as dynamic axis hints

* check dir

* update

* support longformerattention

* update and add support for bert ops

* trim

* review comments

* review comments

onnxruntime/python/tools/transformers/fusion_skiplayernorm.py

@@ -18,6 +18,7 @@ class FusionSkipLayerNormalization(Fusion):
     """
     def __init__(self, model: OnnxModel):
         super().__init__(model, "SkipLayerNormalization", "LayerNormalization")
+        self.shape_infer_helper = self.model.infer_runtime_shape({"batch_size": 4, "seq_len": 7})
 
     def fuse(self, node, input_name_to_nodes, output_name_to_node):
         add = self.model.get_parent(node, 0, output_name_to_node)
@@ -35,6 +36,14 @@ class FusionSkipLayerNormalization(Fusion):
         if len(self.model.get_parents(add)) != 2:
             return
 
+        if self.shape_infer_helper is not None:
+            if not self.shape_infer_helper.compare_shape(add.input[0], add.input[1]):
+                return
+        else:
+            logger.warning(
+                "symbolic shape infer failed. it's safe to ignore this message if there is no issue with optimized model"
+            )
+
         gather_path = self.model.match_parent_path(add, ['Gather'], [None])
         if gather_path is not None and self.model.find_graph_input(gather_path[0].input[1]) is None:
             if self.model.match_parent_path(gather_path[0], ['ConstantOfShape'], [1]) is None:

onnxruntime/python/tools/transformers/onnx_model.py

@@ -12,6 +12,7 @@ from pathlib import Path
 import numpy as np
 from collections import deque
 from onnx import ModelProto, TensorProto, numpy_helper, helper, external_data_helper, save_model
+from shape_infer_helper import SymbolicShapeInferenceHelper
 
 logger = logging.getLogger(__name__)
 
@@ -20,6 +21,21 @@ class OnnxModel:
     def __init__(self, model):
         self.model = model
         self.node_name_counter = {}
+        self.shape_infer_helper = None
+
+    def infer_runtime_shape(self, dynamic_axis_mapping, update = False):
+        shape_infer_helper = None
+        if update:
+            shape_infer_helper = SymbolicShapeInferenceHelper(self.model)
+            self.shape_infer_helper = shape_infer_helper
+        else:
+            if self.shape_infer_helper is None:
+                self.shape_infer_helper = SymbolicShapeInferenceHelper(self.model)
+            shape_infer_helper = self.shape_infer_helper
+
+        if shape_infer_helper.infer(dynamic_axis_mapping):
+            return shape_infer_helper
+        return None
 
     def input_name_to_nodes(self):
         input_name_to_nodes = {}

onnxruntime/python/tools/transformers/shape_infer_helper.py

@@ -0,0 +1,84 @@
+#-------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation.  All rights reserved.
+# Licensed under the MIT License.
+#--------------------------------------------------------------------------
+
+import os
+import sys
+
+# In ORT Package the symbolic_shape_infer.py is in ../tools
+file_path = os.path.dirname(__file__)
+if os.path.exists(os.path.join(file_path, "../tools/symbolic_shape_infer.py")):
+    sys.path.append(os.path.join(file_path, '../tools'))
+else:
+    sys.path.append(os.path.join(file_path, '..'))
+from symbolic_shape_infer import *
+
+
+class SymbolicShapeInferenceHelper(SymbolicShapeInference):
+    def __init__(self, model, verbose=0, int_max=2**31 - 1, auto_merge=True, guess_output_rank=False):
+        super().__init__(int_max, auto_merge, guess_output_rank, verbose)
+        self.model_ = onnx.ModelProto()
+        self.model_.CopyFrom(model)
+        self.all_shapes_inferred_ = False
+        self.inferred_ = False
+
+    # The goal is to remove dynamic_axis_mapping
+    def infer(self, dynamic_axis_mapping):
+        if self.inferred_:
+            return self.all_shapes_inferred_
+
+        self.dynamic_axis_mapping_ = dynamic_axis_mapping  # e.g {"batch_size" : 4, "seq_len" :7}
+
+        self._preprocess(self.model_)
+        while self.run_:
+            self.all_shapes_inferred_ = self._infer_impl()
+
+        self.inferred_ = True
+        return self.all_shapes_inferred_
+
+    # override _preprocess() to avoid unnecessary model copy since ctor copies the model
+    def _preprocess(self, in_mp):
+        self.out_mp_ = in_mp
+        self.initializers_ = dict([(i.name, i) for i in self.out_mp_.graph.initializer])
+        self.known_vi_ = dict([(i.name, i) for i in list(self.out_mp_.graph.input)])
+        self.known_vi_.update(
+            dict([(i.name, helper.make_tensor_value_info(i.name, i.data_type, list(i.dims)))
+                  for i in self.out_mp_.graph.initializer]))
+
+    # Override _get_sympy_shape() in symbolic_shape_infer.py to ensure shape inference by giving the actual value of dynamic axis
+    def _get_sympy_shape(self, node, idx):
+        sympy_shape = []
+        for d in self._get_shape(node, idx):
+            if type(d) == str:
+                if d in self.dynamic_axis_mapping_.keys():
+                    sympy_shape.append(self.dynamic_axis_mapping_[d])
+                elif d in self.symbolic_dims_:
+                    sympy_shape.append(self.symbolic_dims_[d])
+                else:
+                    sympy_shape.append(sympy.Symbol(d, integer=True))
+            else:
+                assert None != d
+                sympy_shape.append(d)
+        return sympy_shape
+
+    def get_edge_shape(self, edge):
+        assert (self.all_shapes_inferred_ == True)
+        if edge not in self.known_vi_:
+            print("Cannot retrive the shape of " + str(edge))
+            return None
+        type_proto = self.known_vi_[edge].type
+        shape = get_shape_from_type_proto(type_proto)
+        for i in range(len(shape)):
+            d = shape[i]
+            if type(d) == str and d in self.dynamic_axis_mapping_.keys():
+                shape[i] = self.dynamic_axis_mapping_[d]
+        return shape
+
+    def compare_shape(self, edge, edge_other):
+        assert (self.all_shapes_inferred_ == True)
+        shape = self.get_edge_shape(edge)
+        shape_other = self.get_edge_shape(edge_other)
+        if shape is None or shape_other is None:
+            raise Exception("At least one shape is missed for edges to compare")
+        return shape == shape_other

onnxruntime/python/tools/transformers/test/test_shape_infer_helper.py

@@ -0,0 +1,48 @@
+import os
+import unittest
+import sys
+sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
+
+from onnx_exporter import export_onnx_model_from_pt
+from huggingface_models import MODELS
+from benchmark_helper import Precision
+from shape_infer_helper import *
+
+
+class SymbolicShapeInferenceHelperTest(unittest.TestCase):
+    def _load_onnx(self, model_name):
+        input_names = MODELS[model_name][0]
+        base_path = "../onnx_models/"
+        import torch
+        with torch.no_grad():
+            export_onnx_model_from_pt(model_name, MODELS[model_name][1], MODELS[model_name][2], MODELS[model_name][3],
+                                      None, '../cache_models', base_path, input_names[:1], False, Precision.FLOAT32,
+                                      True, True, True, False, {})
+        model_path = base_path + model_name.replace('-', '_') + "_1.onnx"
+        import onnx
+        return onnx.load_model(model_path)
+
+    def test_bert_shape_infer_helper(self):
+        model = self._load_onnx("bert-base-cased")
+        shape_infer_helper = SymbolicShapeInferenceHelper(model)
+        self.assertEqual(shape_infer_helper.infer({"batch_size": 4, "seq_len": 16}), True)
+        self.assertEqual(shape_infer_helper.get_edge_shape("802"), [4, 16, 768])
+        self.assertEqual(shape_infer_helper.get_edge_shape("804"), [4, 16, 1])
+        self.assertEqual(shape_infer_helper.get_edge_shape("1748"), [])
+        self.assertEqual(shape_infer_helper.get_edge_shape("encoder.layer.4.attention.output.LayerNorm.weight"), [768])
+        self.assertEqual(shape_infer_helper.get_edge_shape("1749"), [768, 3072])
+        self.assertEqual(shape_infer_helper.get_edge_shape("817"), [4, 16, 3072])
+        self.assertEqual(shape_infer_helper.get_edge_shape("encoder.layer.4.intermediate.dense.bias"), [3072])
+        self.assertEqual(shape_infer_helper.get_edge_shape("1750"), [3072, 768])
+        self.assertEqual(shape_infer_helper.get_edge_shape("853"), [3])
+        self.assertEqual(shape_infer_helper.get_edge_shape("858"), [1])
+        self.assertEqual(shape_infer_helper.get_edge_shape("880"), [4, 16, 12, 64])
+
+        self.assertEqual(shape_infer_helper.compare_shape("329", "253"), True)
+        self.assertEqual(shape_infer_helper.compare_shape("447", "371"), True)
+        self.assertEqual(shape_infer_helper.compare_shape("329", "817"), False)
+        self.assertEqual(shape_infer_helper.compare_shape("447", "853"), False)
+
+
+if __name__ == '__main__':
+    unittest.main()