[英]Tensorflow: delete nodes from graph
I'm trying to delete some nodes from graph and save it in .pb我正在尝试从图中删除一些节点并将其保存在 .pb 中
Only needed nodes can be added to new mod_graph_def
graph, but the problem that graph still have some references to deleted node in other nodes inputs, but I can't modify inputs of node:只有需要的节点才能添加到新的
mod_graph_def
图,但问题是图仍然有一些对其他节点输入中已删除节点的引用,但我无法修改节点的输入:
def delete_ops_from_graph():
with open(input_model_filepath, 'rb') as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
nodes = []
for node in graph_def.node:
if 'Neg' in node.name:
print('Drop', node.name)
else:
nodes.append(node)
mod_graph_def = tf.GraphDef()
mod_graph_def.node.extend(nodes)
# The problem that graph still have some references to deleted node in other nodes inputs
for node in mod_graph_def.node:
inp_names = []
for inp in node.input:
if 'Neg' in inp:
pass
else:
inp_names.append(inp)
node.input = inp_names # TypeError: Can't set composite field
with open(output_model_filepath, 'wb') as f:
f.write(mod_graph_def.SerializeToString())
def delete_ops_from_graph():
with open(input_model_filepath, 'rb') as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
# Delete nodes
nodes = []
for node in graph_def.node:
if 'Neg' in node.name:
print('Drop', node.name)
else:
nodes.append(node)
mod_graph_def = tf.GraphDef()
mod_graph_def.node.extend(nodes)
# Delete references to deleted nodes
for node in mod_graph_def.node:
inp_names = []
for inp in node.input:
if 'Neg' in inp:
pass
else:
inp_names.append(inp)
del node.input[:]
node.input.extend(inp_names)
with open(output_model_filepath, 'wb') as f:
f.write(mod_graph_def.SerializeToString())
Previous answer is good but I would suggest to tie removed node input to the next node input.以前的答案很好,但我建议将删除的节点输入与下一个节点输入联系起来。 Like if we have a chain
A-input b->B-input c->C-input d->D
and going to remove say node B
then we should not just remove input c
but replace it with input b
.就像如果我们有一个链
A-input b->B-input c->C-input d->D
并且要删除节点B
那么我们不应该只是删除input c
而是用input b
替换它。 Look at code below:看下面的代码:
# remove node and connect its input to follower
def remove_node(graph_def, node_name, input_name):
nodes = []
for node in graph_def.node:
if node.name == node_name:
assert(input_name in node.input or len(node.input) == 0),\
"Node input to use is not among inputs of node to remove"
input_of_removed_node = input_name if len(node.input) else ''
print("Removing {} and using its input {}".format(node.name,
input_of_removed_node))
continue
nodes.append(node)
# modify inputs where required
# removed name must be replaced with input from removed node
for node in nodes:
inp_names = []
replace = False
for inp in node.input:
if inp == node_name:
inp_names.append(input_of_removed_node)
print("For node {} replacing input {}
with {}".format(node.name, inp, input_of_removed_node))
replace = True
else:
inp_names.append(inp)
if replace:
del node.input[:]
node.input.extend(inp_names)
mod_graph_def = tf.GraphDef()
mod_graph_def.node.extend(nodes)
return mod_graph_def
声明:本站的技术帖子网页,遵循CC BY-SA 4.0协议,如果您需要转载,请注明本站网址或者原文地址。任何问题请咨询:yoyou2525@163.com.