tf.data 内存泄漏 [英] Memory leak with tf.data

问题描述

我在 for 循环中创建了一个 tf.data.Dataset，我注意到内存并没有像预期的那样在每次迭代后释放.

I'm creating a tf.data.Dataset inside a for loop and I noticed that the memory was not freed as one would expect after each iteration.

有没有办法从 TensorFlow 请求释放内存?

Is there a way to request from TensorFlow to free the memory?

我尝试使用 tf.reset_default_graph()，我尝试在相关的 python 对象上调用 del 但这不起作用.

I tried using tf.reset_default_graph(), I tried calling del on the relevant python objects but this does not work.

似乎唯一有效的是 gc.collect().不幸的是，gc.collect 不适用于一些更复杂的示例.

The only thing that seems to work is gc.collect(). Unfortunately, gc.collect does not work on some more complex examples.

完全可重现的代码:

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
import psutil
%matplotlib inline

memory_used = []
for i in range(500):
    data = tf.data.Dataset.from_tensor_slices(
                    np.random.uniform(size=(10, 500, 500)))\
                    .prefetch(64)\
                    .repeat(-1)\
                    .batch(3)
    data_it = data.make_initializable_iterator()
    next_element = data_it.get_next()

    with tf.Session() as sess:
        sess.run(data_it.initializer)
        sess.run(next_element)
    memory_used.append(psutil.virtual_memory().used / 2 ** 30)
    tf.reset_default_graph()

plt.plot(memory_used)
plt.title('Evolution of memory')
plt.xlabel('iteration')
plt.ylabel('memory used (GB)')

推荐答案

此修复 当我在 TF 2.4 中遇到类似问题时为我工作

This fix worked for me when I had a similar issue with TF 2.4

1. sudo apt-get install libtcmalloc-minimal4
2. LD_PRELOAD=/path/to/libtcmalloc_minimal.so.4 python example.py

这篇关于tf.data 内存泄漏的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！