Tensorflow在每次调用带有最终图形的session.run()时泄漏内存

时间:2019-05-14 04:50:28

标签: python tensorflow memory-leaks tensorflow-datasets

我正在尝试使用tf.data API将可变大小的图像数据(LxLx2)馈入模型,但是我注意到每次迭代都会泄漏内存。我希望内存使用量将由数据集中最大的图像确定,但是我可以看到,即使处理的图像小于到目前为止所看到的最大大小,内存使用量也在增加。

Leaking memory over 100 iterations

当我直接收集处理后的特征而不是计算神经网络激活时,内存似乎没有泄漏。

Expected memory use (forgoing NN computation)

这种类型问题的最常见原因似乎是在图上动态添加节点,但是我在迭代之前调用了graph.finalize()并且没有捕获任何错误。

我正在使用python 3.5.4和tensorflow 1.10并仅在CPU上运行计算。

import tensorflow as tf
from sys import argv

# Data preparation
def record_parser(value):
    keys_to_features = {
        'seq_length': tf.VarLenFeature(dtype=tf.int64),
        'seq_feat': tf.VarLenFeature(dtype=tf.float32)
        }
    parsed = tf.parse_single_example(value, keys_to_features)
    length_ = tf.reshape(parsed['seq_length'].values, [])
    i32_len = tf.cast(length_, dtype=tf.int32)
    features_ = tf.reshape(parsed['seq_feat'].values, [i32_len, i32_len, 2])
    return features_

graph = tf.get_default_graph()
dataset_ = tf.data.TFRecordDataset(argv[1])
dataset_ = dataset_.map(lambda value: record_parser(value))
dataset_ = dataset_.batch(1)
iterator = dataset_.make_one_shot_iterator()
features = iterator.get_next()

# NN part
nn0 = tf.layers.conv2d(features, filters=64, kernel_size=15, padding='SAME',\
 activation=tf.nn.relu)
nn = tf.layers.dense(nn0, units=100, activation=tf.nn.relu)
prediction = tf.layers.dense(nn, 17, activation=None)

var_init_op = tf.group(
                tf.global_variables_initializer(),
                tf.local_variables_initializer()
                )
graph.finalize()

# Iterating over samples
with tf.Session() as sess:
    sess.run(var_init_op)
    for i in range(100):
        out_loss = sess.run(prediction)
        #out_loss = sess.run(features)

1 个答案:

答案 0 :(得分:0)

在此提及答案,以造福社区。

tf.data中使用Tensorflow Version 1.10 API时的内存泄漏问题已通过升级到Tensorflow Version 1.13解决。