From 79097ef5535cc5ac18fc8e9010c99de08df21340 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Fri, 23 Jul 2021 18:12:21 -0700 Subject: [PATCH] remove useless reshape node (#8419) --- .../python/tools/transformers/fusion_utils.py | 22 +++++++++++++++++++ .../tools/transformers/onnx_model_bert.py | 3 +++ 2 files changed, 25 insertions(+) diff --git a/onnxruntime/python/tools/transformers/fusion_utils.py b/onnxruntime/python/tools/transformers/fusion_utils.py index f491cace14..600f1b6df7 100644 --- a/onnxruntime/python/tools/transformers/fusion_utils.py +++ b/onnxruntime/python/tools/transformers/fusion_utils.py @@ -100,6 +100,28 @@ class FusionUtils: else: return value == expected_value + @staticmethod + def remove_useless_reshape_nodes(model:OnnxModel): + """Remove reshape node that is not needed based on symbolic shape inference: input and output has same shape + """ + shape_infer = model.infer_runtime_shape(update=True) + if shape_infer is None: + return + + nodes_to_remove = [] + for node in model.nodes(): + if node.op_type == 'Reshape': + input_shape = shape_infer.get_edge_shape(node.input[0]) + output_shape = shape_infer.get_edge_shape(node.output[0]) + if input_shape and output_shape and input_shape == output_shape: + logger.info(f"Remove reshape node {node.name} since its input shape is same as output: {input_shape}") + nodes_to_remove.append(node) + + if nodes_to_remove: + for node in nodes_to_remove: + model.replace_input_of_all_nodes(node.output[0], node.input[0]) + model.remove_node(node) + model.prune_graph() class NumpyHelper: @staticmethod diff --git a/onnxruntime/python/tools/transformers/onnx_model_bert.py b/onnxruntime/python/tools/transformers/onnx_model_bert.py index 3287daf410..20e33687b0 100644 --- a/onnxruntime/python/tools/transformers/onnx_model_bert.py +++ b/onnxruntime/python/tools/transformers/onnx_model_bert.py @@ -185,6 +185,9 @@ class BertOnnxModel(OnnxModel): return def adjust_reshape_and_expand(self): + # Remove reshape nodes that having same shape of input and output based on symbolic shape inference. + FusionUtils.remove_useless_reshape_nodes(self) + nodes_to_remove = [] for node in self.nodes(): if node.op_type == 'Reshape':