[英]Tensorflow: delete nodes from graph
我正在嘗試從圖中刪除一些節點並將其保存在 .pb 中
只有需要的節點才能添加到新的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())
以前的答案很好,但我建議將刪除的節點輸入與下一個節點輸入聯系起來。 就像如果我們有一個鏈A-input b->B-input c->C-input d->D
並且要刪除節點B
那么我們不應該只是刪除input c
而是用input b
替換它。 看下面的代碼:
# 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.