saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-19 19:00:47 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

remove useless reshape node (#8419)

Browse source
This commit is contained in:
Tianlei Wu 2021-07-23 18:12:21 -07:00 committed by GitHub
parent 6dee9b9d2d
commit 79097ef553
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 25 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

22
onnxruntime/python/tools/transformers/fusion_utils.py
View file

@ -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

3
onnxruntime/python/tools/transformers/onnx_model_bert.py
View file

@ -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':