时间戳

首页 > 解决方案 > 在循环中构建图形时Tensorflow内存泄漏

python - 在循环中构建图形时Tensorflow内存泄漏

问题描述

当我的网格搜索选择 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。但是不断的增加仍然存在，只是变慢了。

标签： pythonloopstensorflowmemory-leaks

解决方案

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

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()

推荐阅读