删除 TensorFlow 图中除少数节点以外的所有节点

Deleting all but a few nodes in TensorFlow graph

我的 TensorFlow 用例要求我为每个需要处理的实例构建一个新的计算图。这最终会增加内存需求。

除了少数 tf.Variables 是模型参数,我想删除所有其他节点。其他有类似问题的人发现 tf.reset_default_graph() 很有用,但这会去掉我需要保留的模型参数。

我可以使用什么来删除除这些节点以外的所有节点?

编辑: 特定于实例的计算实际上只是意味着我要添加很多新操作。我相信这些操作是内存问题背后的原因。

更新: 请参阅最近发布的 tensorflow fold (https://github.com/tensorflow/fold),它允许动态构建计算图。

tf.graph 数据结构被设计为 append-only 数据结构。因此无法删除或修改现有节点。通常这不是问题,因为在 运行 会话时只处理必要的子图。

您可以尝试将图表的变量复制到新图表中并删除旧图表。要存档这只是 运行:

old_graph = tf.get_default_graph() # Save the old graph for later iteration
new_graph = tf.graph() # Create an empty graph
new_graph.set_default() # Makes the new graph default

如果要遍历旧图中的所有节点,请使用:

for node in old_graph.get_operations():
    if node.type == 'Variable':
       # read value of variable and copy it into new Graph

或者您可以使用:

for node in old_graph.get_collection('trainable_variables'):
   # iterates over all trainable Variabels
   # read and create new variable

另请参阅 python/framework/ops.py : 1759 以了解更多操作图中节点的方法。

然而,在你乱搞 tf.Graph 之前,我强烈建议你考虑一下这是否真的需要。通常可以尝试概括计算并使用共享变量构建一个图,这样你要处理的每个实例都是这个图的子图。