mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-30 20:18:08 +00:00
fix duplicated node name (#7865)
This commit is contained in:
parent
1f4421fe70
commit
71b05f74a2
3 changed files with 30 additions and 17 deletions
|
|
@ -7,7 +7,6 @@ from pathlib import Path
|
|||
class ONNXModel:
|
||||
def __init__(self, model):
|
||||
self.model = model
|
||||
self.node_name_counter = {}
|
||||
|
||||
def nodes(self):
|
||||
return self.model.graph.node
|
||||
|
|
|
|||
|
|
@ -112,7 +112,7 @@ class FusionLayerNormalization(Fusion):
|
|||
inputs=[node.input[0], weight_input, bias_input],
|
||||
outputs=[last_add_node.output[0]],
|
||||
name=self.model.create_node_name("LayerNormalization",
|
||||
name_prefix="SkipLayerNorm"))
|
||||
name_prefix="LayerNorm"))
|
||||
normalize_node.attribute.extend([helper.make_attribute("epsilon", float(add_weight))])
|
||||
self.nodes_to_add.append(normalize_node)
|
||||
self.node_name_to_graph_name[normalize_node.name] = self.this_graph_name
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@
|
|||
# Licensed under the MIT License.
|
||||
#--------------------------------------------------------------------------
|
||||
|
||||
from typing import List, Tuple
|
||||
from typing import List, Tuple, Dict
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
|
|
@ -19,7 +19,7 @@ logger = logging.getLogger(__name__)
|
|||
class OnnxModel:
|
||||
def __init__(self, model):
|
||||
self.model = model
|
||||
self.node_name_counter = {}
|
||||
self._node_name_suffix: Dict[str, int] = {} # key is node name prefix, value is the last suffix generated
|
||||
self.shape_infer_helper = None
|
||||
self.all_graphs = None
|
||||
|
||||
|
|
@ -553,25 +553,39 @@ class OnnxModel:
|
|||
cast_node.attribute.extend([helper.make_attribute("to", int(TensorProto.FLOAT))])
|
||||
self.add_node(cast_node)
|
||||
|
||||
# create a new name for node
|
||||
def create_node_name(self, op_type, name_prefix=None):
|
||||
if op_type in self.node_name_counter:
|
||||
self.node_name_counter[op_type] += 1
|
||||
"""Create a unique node name that starts with a prefix (default is operator type).
|
||||
The name will not be duplicated with any name that generated or existed in current graphs.
|
||||
Args:
|
||||
op_type (str): operator type
|
||||
name_prefix (str, optional): prefix of node name. Defaults to None.
|
||||
|
||||
Returns:
|
||||
str: node name
|
||||
"""
|
||||
|
||||
if name_prefix:
|
||||
prefix = name_prefix if name_prefix.endswith("_") else (name_prefix + "_")
|
||||
else:
|
||||
self.node_name_counter[op_type] = 1
|
||||
prefix = op_type + "_"
|
||||
|
||||
if name_prefix is not None:
|
||||
full_name = name_prefix + str(self.node_name_counter[op_type])
|
||||
suffix: int = 0
|
||||
if prefix in self._node_name_suffix:
|
||||
suffix = self._node_name_suffix[prefix] + 1
|
||||
else:
|
||||
full_name = op_type + "_" + str(self.node_name_counter[op_type])
|
||||
# Check existed node name only once for a prefix as we assume create_node_name is called for every new node in fusion.
|
||||
for node in self.nodes():
|
||||
if node.name and node.name.startswith(prefix):
|
||||
try:
|
||||
index = int(node.name[len(prefix):])
|
||||
suffix = max(index + 1, suffix)
|
||||
except ValueError:
|
||||
continue
|
||||
|
||||
# Check whether the name is taken:
|
||||
nodes = self.get_nodes_by_op_type(op_type)
|
||||
for node in nodes:
|
||||
if node.name == full_name:
|
||||
raise Exception("Node name already taken:", full_name)
|
||||
# Record the generated suffix so that we can avoid generating duplicated name.
|
||||
self._node_name_suffix[prefix] = suffix
|
||||
|
||||
return full_name
|
||||
return prefix + str(suffix)
|
||||
|
||||
def find_graph_input(self, input_name):
|
||||
for input in self.model.graph.input:
|
||||
|
|
|
|||
Loading…
Reference in a new issue