From 2ad94d7042aefc324cb854d9eac6ce3ab8bfb3f3 Mon Sep 17 00:00:00 2001 From: ZhangGe6 Date: Sat, 11 Jun 2022 17:50:26 +0800 Subject: [PATCH] fix custom added node deleting issues --- onnx_modifier.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/onnx_modifier.py b/onnx_modifier.py index 5eb6d03..20228c1 100644 --- a/onnx_modifier.py +++ b/onnx_modifier.py @@ -102,8 +102,10 @@ class onnxModifier: node.output[i] = dst_name - def add_node(self, nodes_info): + def add_node(self, nodes_info, node_states): for node_info in nodes_info.values(): + if node_states[node_info['properties']['name']] == "Deleted": + continue node = make_node(node_info) # print(node) @@ -115,7 +117,7 @@ class onnxModifier: print(modify_info['added_node_info']) self.remove_node_by_node_states(modify_info['node_states']) self.modify_node_io_name(modify_info['node_renamed_io']) - self.add_node(modify_info['added_node_info']) + self.add_node(modify_info['added_node_info'], modify_info['node_states']) def check_and_save_model(self, save_dir='./modified_onnx'): if not os.path.exists(save_dir):