为什么这个数据集实现 运行 内存不足?

Why would this dataset implementation run out of memory?

我关注this instruction and write the following code to create a Dataset for images(COCO2014 training set)

from pathlib import Path
import tensorflow as tf


def image_dataset(filepath, image_size, batch_size, norm=True):
    def preprocess_image(image):
        image = tf.image.decode_jpeg(image, channels=3)
        image = tf.image.resize(image, image_size)
        if norm:
            image /= 255.0  # normalize to [0,1] range
        return image

    def load_and_preprocess_image(path):
        image = tf.read_file(path)
        return preprocess_image(image)

    all_image_paths = [str(f) for f in Path(filepath).glob('*')]
    path_ds = tf.data.Dataset.from_tensor_slices(all_image_paths)
    ds = path_ds.map(load_and_preprocess_image, num_parallel_calls=tf.data.experimental.AUTOTUNE)
    ds = ds.shuffle(buffer_size = len(all_image_paths))
    ds = ds.repeat()
    ds = ds.batch(batch_size)
    ds = ds.prefetch(tf.data.experimental.AUTOTUNE)

    return ds

ds = image_dataset(train2014_dir, (256, 256), 4, False)
image = ds.make_one_shot_iterator().get_next('images')
# image is then fed to the network

此代码将始终 运行 内存 (32G) 和 GPU (11G) 并终止进程。这是终端上显示的消息。

我还发现程序卡在 sess.run(opt_op)。哪里错了?我该如何解决?

问题是这样的:

ds = ds.shuffle(buffer_size = len(all_image_paths))

Dataset.shuffle() 使用的缓冲区是 'in memory' 缓冲区,因此您实际上是在尝试将整个数据集加载到内存中。

您有几个选项(可以结合使用)来解决此问题:

选项 1:

将缓冲区大小减小到更小的数量。

选项 2:

shuffle() 语句移到 map() 语句之前。

这意味着我们将在加载图像之前洗牌,因此我们只是将文件名存储在内存缓冲区中以进行洗牌,而不是存储巨大的张量。