python - tf.data.Dataset - 为什么当我缓存示例时数据管道的性能没有提高？

Question

我目前正在尝试了解有关使用 tf.data 构建高效预处理管道的更多信息。根据this tutorial缓存数据时对性能的影响应该是不可忽略的。

我将数据管道简化为一个非常简单的示例来验证这种效果。

import os
import tensorflow as tf

class ExperimentalDS:
    def __init__(self, hr_img_path, cache, repeat, shuffle_buffer_size=4096):
        self.hr_img_path = hr_img_path
        self.ids = os.listdir(self.hr_img_path)
        self.train_list = self.ids

        train_list_ds = tf.data.Dataset.list_files([f"{hr_img_path}/{fname}" for fname in self.train_list])

        train_hr_ds = train_list_ds.map(self.load_img)
        train_hr_ds = train_hr_ds.shuffle(shuffle_buffer_size)

        self.train_ds = train_hr_ds

        # should probably call shuffle again after caching
        if cache: self.train_ds.cache()
        self.train_ds = train_hr_ds.repeat(repeat)

    def get_train_ds(self, batch_size=8):
        return self.train_ds.batch(batch_size).prefetch(tf.data.experimental.AUTOTUNE)

    def load_img(self, fpath):
        img = tf.io.read_file(fpath)
        img = tf.image.decode_png(img)
        img = tf.image.convert_image_dtype(img, tf.float32)
        return img

管道基本上只是从文件夹中读取文件名，从这些文件名加载图像，随机排列图像，然后缓存它们或不缓存它们，具体取决于提供的参数。

为了评估性能，我主要复制了前面提到的教程中的基准测试函数。

def benchmark_dataset(ds, num_steps):
    start = time.perf_counter()
    it = iter(ds)

    for i in range(num_steps):
        batch = next(it)
        if i % 100 == 0:
            print(".", end="")
    print()

    end = time.perf_counter()
    duration = end - start

    return duration

if __name__ == "__main__":
    num_steps = 1000
    batch_size = 8
    durations_no_cache = []
    durations_cached = []
    for i in range(num_steps):
        ds = ExperimentalDS("./test_data/benchmark/16", cache=False, repeat=-1)
        ds_train = ds.get_train_ds(batch_size=batch_size)
        durations_no_cache.append(benchmark_dataset(ds_train, num_steps))

    for i in range(num_steps):
        ds = ExperimentalDS("./test_data/benchmark/16", cache=True, repeat=-1)
        ds_train = ds.get_train_ds(batch_size=batch_size)
        durations_cached.append(benchmark_dataset(ds_train, num_steps))

    os.makedirs(SAVE_PATH, exist_ok=True)
    durations_no_cache = np.array(durations_no_cache)
    avg_duration_no_cache = np.average(durations_no_cache)

    durations_cached = np.array(durations_cached)
    avg_durations_cached = np.average(durations_cached)

    with open(f"{SAVE_PATH}/stats", "a+") as f:
        f.write("no cache:\n")
        f.write(f"{num_steps} batches: {avg_duration_no_cache}s (avg)\n")
        f.write(f"{batch_size*num_steps/avg_duration_no_cache:.5f} Images/s\n\n")
        f.write("cached:\n")
        f.write(f"{num_steps} batches: {avg_durations_cached}s (avg)\n")
        f.write(f"{batch_size*num_steps/avg_durations_cached:.5f} Images/s")

我正在加载一个非常简单的图像数据集，其中包含 16 个图像，每个图像的尺寸为 128x128(因此它应该很容易装入内存)。我无限期地重复此数据集，并在使用缓存和不使用缓存的情况下迭代 1000 个批处理(批处理大小为 8)，记录运行时间，然后对 1000 次运行的结果进行平均。由于这些运行次数相当多，因此我认为应该不会有太大差异。如果重要的话，基准测试是在 GPU 上运行的。

结果令我非常惊讶。没有缓存的基准测试实际上稍微快一些:

no cache:
1000 batches: 2.434403038507444s (avg)
3286.22659 Images/s

cached:
1000 batches: 2.439824645938235s (avg)
3278.92417 Images/s

我知道还有一些其他方法可以提高性能，例如并行和矢量化映射，但它在比较缓存与不缓存方面不应该产生任何影响。

有人可以帮我解决这个问题吗？我在这里缺少什么？

编辑:在评论中，@Szymon Maszke 建议我应该对多个时期的迭代进行基准测试，并将数据实际馈送到网络。所以我这样做了，但是缓存和未缓存的数据集的性能几乎相同。真的不知道为什么。

edit2:修复了@AAudibert指出的错误后，它现在按预期工作。事实上，老实说，它的工作效果比预期的要好:

no cache:
1000 batches: 2.624478972374927s (avg)
3048.22408 Images/s

cached:
1000 batches: 0.17946020061383025s (avg)
44578.12915 Images/s

Answer 1

此语句不执行任何操作:

if cache: self.train_ds.cache()

应该是:

if cache: train_hr_ds = train_hr_ds.cache()

与其他数据集转换一样，缓存返回新数据集，而不是修改现有数据集。