在Tensorflow中进行批量训练时如何限制RAM使用量？

Question

我正在使用大小为40的小批量的大型图像数据集训练一个深度神经网络 。我的数据集为.mat格式（如果需要，我可以轻松将其更改为任何其他格式，例如.npy格式），并且在进行训练之前，加载为一个4-D numpy数组。 我的问题是，在训练时，cpu-RAM（不是GPU RAM）很快就会耗尽，并开始使用我的Swap内存的几乎一半。

我的训练代码具有以下模式：

batch_size = 40
...
with h5py.File('traindata.mat', 'r') as _data:
    train_imgs = np.array(_data['train_imgs'])

# I can replace above with below loading, if necessary
# train_imgs = np.load('traindata.npy')

...

shape_4d = train_imgs.shape 
for epoch_i in range(max_epochs):
    for iter in range(shape_4d[0] // batch_size):
        y_ = train_imgs[iter*batch_size:(iter+1)*batch_size]
        ...
        ...

似乎完整训练数据的初始加载本身已成为瓶颈（在我中止之前接管了12 GB cpu RAM）。

解决这个瓶颈的最佳有效方法是什么？

提前致谢。

Answer 1

在内存中加载大数据集不是一个好主意。 我建议您使用不同的方式加载数据集，看看TensorFlow中的数据集API： https ：//www.tensorflow.org/programmers_guide/datasets

您可能需要将数据转换为其他格式，但是如果您有一个CSV或TXT文件（每行带有一个示例），则可以使用TextLineDataset并向其提供模型：

filenames = ["/var/data/file1.txt", "/var/data/file2.txt"]
dataset = tf.data.TextLineDataset(filenames)

def _parse_py_fun(text_line):
    ... your custom code here, return np arrays

def _map_fun(text_line):
    result = tf.py_func(_parse_py_fun, [text_line], [tf.uint8])
    ... other tensorlow code here
    return result

dataset = dataset.map(_map_fun)
dataset = dataset.batch(4)
iterator = dataset.make_one_shot_iterator()
input_data_of_your_model = iterator.get_next()

output = build_model_fn(input_data_of_your_model)

sess.run([output]) # the input was assigned directly when creating the model