在循环中构建图时，Tensorflow内存泄漏

Question

当我因选择消耗大量Tensorflow（版本1.12.0）模型的网格而崩溃时，我注意到了这一点。

请注意，与这里的类似问题不同，我确实关闭了图和会话（使用上下文管理器），并且没有在循环中向图添加节点。

我怀疑tensorflow可能保留了在每次迭代之间都不会清除的全局变量，因此我在迭代之前和之后都调用了globals（），但是在每次迭代之前和之后都没有观察到全局变量集合中的任何差异。 >

我举了一个小例子来重现问题。我在循环中训练了一个简单的MNIST分类器，并绘制了进程消耗的内存：

import matplotlib.pyplot as plt
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
import psutil
import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data
process = psutil.Process(os.getpid())

N_REPS = 100
N_ITER = 10
mnist = input_data.read_data_sets("MNIST_data/", one_hot=True)
x_test, y_test = mnist.test.images, mnist.test.labels

# Runs experiment several times.
mem = []
for i in range(N_REPS):
    with tf.Graph().as_default():
        net = tf.contrib.layers.fully_connected(x_test, 200)
        logits = tf.contrib.layers.fully_connected(net, 10, activation_fn=None)
        loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(labels=y_test, logits=logits))
        train_op = tf.train.AdamOptimizer(learning_rate=0.0001).minimize(loss)
        init = tf.global_variables_initializer()
        with tf.Session() as sess:
            # training loop.
            sess.run(init)
            for _ in range(N_ITER):
                sess.run(train_op)
    mem.append(process.memory_info().rss)
plt.plot(range(N_REPS), mem)

结果图如下：

在我的实际项目中，进程内存从几百MB（取决于数据集的大小）开始，一直到64 GB，直到我的系统内存不足为止。我尝试过一些使增速变慢的方法，例如使用占位符和feed_dicts而不是依赖convert_to_tensor。但是持续的增长仍然存在，只是速度较慢。

Answer 1

您需要在for循环的每次迭代之后清除图形，然后实例化新图形。在循环的末尾添加tf.reset_default_graph()应该可以解决您的内存泄漏问题。

for i in range(N_REPS):
    with tf.Graph().as_default():
        net = tf.contrib.layers.fully_connected(x_test, 200)
        ...
    mem.append(process.memory_info().rss)
    tf.reset_default_graph()

Answer 2

尝试在会话内部进行循环。不要为每次迭代创建图和会话。每次创建图并初始化变量时，您不是在重新定义旧图，而是创建新图导致内存泄漏。我遇到了类似的问题，并且能够通过在会话中进行循环来解决此问题。

来自How not program Tensorflow

  

  
• 当您创建操作时，请注意，仅创建所需的操作。尝试使操作创建与操作执行保持不同。
  
• 尤其是如果您仅使用默认图形并在常规REPL或笔记本中交互运行，则图形中可能会出现很多废弃的操作。每次重新运行定义任何图形操作的笔记本单元时，您不仅在重新定义操作，还创建了新的操作。