为什么在使用Tensorflow加载批次进行训练时为什么会出现内存泄漏？

Question

I am training a neural network using batches of images. I want to import the images in the for training loop to avoid importing all the images at once. Once a loop iteration is done, I want to use a new batch and forget about previous images to free the CPU memory. I am reallocating the variables every time but the memory (and thus the running time) keeps on increasing. Do you know how to free memory from the previous batches?

I am using Python 3.6.8 and Tensorflow 1.14.0 on a GPU Tesla K80 (memory_limit:11.3 GB).

I have tried gc.collect() but it does not work.

import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt
import os
import time
import gc
import psutil

dir_img = "../dir_img_png/data/"
data = [os.path.join(dir_img, f) for f in os.listdir(dir_img) if os.path.isfile(os.path.join(dir_img, f))]
np.random.shuffle(data)
data_train = data[10:]
data_test = data[:10]

# Hyperparameters
input_size=256
batch_size=64
epochs = 20

def memory():
    pid = os.getpid()
    py = psutil.Process(pid)
    memoryUse = py.memory_info()[0]/2.**30  # memory use in GB
    print('memory use:', memoryUse)

def preprocess_image(path):
    raw_img = tf.read_file(path)
    img = tf.io.decode_png(raw_img, channels=1)
    img = tf.cast(img, tf.float32)
    img -= 127.5
    img *= 1./127.5
    return img

# Networks IO
real_images = tf.placeholder(tf.float32, [None, input_size, input_size, 1])

init = tf.global_variables_initializer()

sess = tf.Session()
sess.run(init)

for epoch in range(epochs):
    for batch in range(len(data_train)//batch_size):
        tps = time.time()
        batch_images = np.array([sess.run(preprocess_image(d)) for d in data_train[batch*batch_size:batch*batch_size+batch_size]])
        print("tps1: {t}".format(t=time.time()-tps))

        tps2 = time.time()
        gc.collect()
        print("tps2: {t}".format(t=time.time()-tps2))

        memory()

sess.close()

Here is the output I get:

tps1: 10.445663928985596
tps2: 0.0995786190032959
memory use: 0.871917724609375
tps1: 9.142687320709229
tps2: 0.10767912864685059
memory use: 0.9062271118164062
tps1: 12.030094146728516
tps2: 0.10630679130554199
memory use: 0.9415740966796875
tps1: 13.415296077728271
tps2: 0.11185669898986816
memory use: 0.9608650207519531
tps1: 12.053950548171997
tps2: 0.11706066131591797
memory use: 0.9794692993164062
tps1: 14.279714584350586
tps2: 0.11610865592956543
memory use: 0.9980583190917969
tps1: 11.772900342941284
tps2: 0.12384176254272461
memory use: 1.0166587829589844
tps1: 15.43606686592102
tps2: 0.12571096420288086

The memory and running time keep on increasing.

Answer 1

Every method tf.* (like tf.read_file, tf.io.decode_png or tf.cast) must be outside of the loop.

We need to define the graph before the loop and then run only the tensor we want :

import numpy as np
import tensorflow as tf
import os
import time
import psutil

dir_img = "../dir_img_png/data/"
data = [os.path.join(dir_img, f) for f in os.listdir(dir_img) if os.path.isfile(os.path.join(dir_img, f))]
np.random.shuffle(data)
data_train = data[10:]
data_test = data[:10]

# Hyperparameters
input_size=256
batch_size=64
epochs = 20

def memory():
    pid = os.getpid()
    py = psutil.Process(pid)
    memoryUse = py.memory_info()[0]/2.**30  # memory use in GB
    print('memory use:', memoryUse)

# Networks IO
img_path = tf.placeholder(tf.string)
real_images = tf.placeholder(tf.float32, [None, input_size, input_size, 1])

# Preprocess image
raw_img = tf.read_file(img_path)
png_img = tf.io.decode_png(raw_img, channels=1)
png_img = tf.cast(png_img, tf.float32)
png_img = tf.subtract(png_img, 127.5)
png_img = tf.divide(png_img, 127.5)

init = tf.global_variables_initializer()

sess = tf.Session()
sess.run(init)

for epoch in range(epochs):
    for batch in range(len(data_train)//batch_size):
        tps = time.time()
        batch_images = np.array([sess.run(png_img, feed_dict={img_path: d}) for d in data_train[batch*batch_size:batch*batch_size+batch_size]])
        print("tps1: {t}".format(t=time.time()-tps))

        memory()

sess.close()

The output shows then that the memory is constant:

tps1: 11.155524492263794
memory use: 0.8702735900878906
tps1: 9.429716110229492
memory use: 0.8859291076660156
tps1: 9.62732195854187
memory use: 0.8859291076660156
tps1: 11.327840089797974
memory use: 0.8859291076660156
tps1: 8.580215215682983
memory use: 0.8859291076660156
tps1: 11.039035081863403
memory use: 0.8859291076660156
tps1: 12.866767168045044
memory use: 0.8859291076660156
tps1: 12.425249576568604
memory use: 0.8859291076660156