Утечка памяти Tensorflow при построении графа в цикле

Question

Я заметил это, когда мой поиск по сетке для выбора гиперпараметров модели Tensorflow (версия 1.12.0) потерпел крах из-за взрыва в потреблении памяти.

Обратите внимание, что в отличие от подобного вопроса здесь, я делаюзакройте график и сессию (используя менеджеры контекста), и я не добавляю узлы к графику в цикле.

Я подозревал, что, возможно, тензор потока поддерживает глобальные переменные, которые не очищаются между итерациями, поэтому я назвал глобальные переменные() до и после итерации, но не наблюдалось никакой разницы в наборе глобальной переменной до и после каждой итерации.

Я сделал небольшой пример, который воспроизводит проблему.Я тренирую простой классификатор 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)

И результирующий график выглядит так:

В моем реальном проекте объем памяти процесса начинается с нескольких сотен МБ (в зависимости от размера набора данных) и достигает 64 ГБ, пока в моей системе не будет исчерпан объем памяти.Есть вещи, которые я пробовал замедлить увеличение, например, использование заполнителей и feed_dicts вместо того, чтобы полагаться на convert_to_tensor.Но постоянный рост все еще там, только медленнее.

Answer 1

Попробуйте взять цикл внутри сессии.Не создавайте график и сессию для каждой итерации.Каждый раз, когда граф создается и инициализируется переменной, вы не переопределяете старый граф, а создаете новый, что приводит к утечкам памяти.Я столкнулся с подобной проблемой и смог решить ее, взяв цикл внутри сеанса.

С Как не программировать Tensorflow

• Beосознавая, когда вы создаете операции, и создавайте только те, которые вам нужны.Постарайтесь не создавать операции отличные от операции.
• Особенно, если вы просто работаете с графиком по умолчанию и работаете в интерактивном режиме в обычном REPL или записной книжке, вы можете получитьмного заброшенных операций в вашем графике.Каждый раз, когда вы перезапускаете ячейку ноутбука, в которой определены какие-либо графические операции, вы не просто переопределяете операции - вы создаете новые.

Answer 2

Вам необходимо очищать график после каждой итерации цикла 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()