tf.data.Dataset - Почему производительность моего конвейера данных не увеличивается, когда я кеширую примеры?

Question

В настоящее время я пытаюсь узнать больше о создании эффективных конвейеров предварительной обработки с помощью tf.data. * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * Нет, способ, что Кэширование данных должно оказывать незначительное влияние на производительность.

Я сократил конвейер данных до очень простого примера для проверки этого эффекта.

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 прогонов. Поскольку это довольно много прогонов, я бы предположил, что не должно быть большой разницы. Тест производительности выполнялся на графическом процессоре, если он имеет значение.

Результаты меня очень удивляют. Тест без кэширования на самом деле немного быстрее:

без кэша:
1000 пакетов: 2.434403038507444s (avg)
3286.22659 Изображения / с

кэширование:
1000 пакетов: 2.439824645938235s (avg)
3278.92417 изображений / с

Я знаю, что есть несколько других вещей, которые могут улучшить производительность, например, параллельное и векторизованное отображение, но оно не должно иметь каких-либо эффект в отношении сравнения кэширования и не кэширования.

Может ли кто-нибудь помочь мне в этом? Что мне здесь не хватает?

edit: в комментариях @Szymon Maszke было предложено сравнить итерацию по нескольким эпохам и фактически передать данные в сеть. Я так и сделал, но набор данных в кэше и без кэша работает примерно одинаково. На самом деле не уверен, почему.

edit2: После исправления ошибки, указанной @AAudibert, она работает, как и ожидалось. На самом деле это работает лучше, чем ожидалось, если честно:

без кеша:
1000 пакетов: 2.624478972374927s (avg)
3048.22408 Изображения / с

кэшировано:
1000 пакетов: 0,17946020061383025s (в среднем)
44578.12915 Изображения / с

Answer 1

Это утверждение ничего не делает:

if cache: self.train_ds.cache()

Оно должно быть:

if cache: train_hr_ds = train_hr_ds.cache()

Как и другие преобразования набора данных, cache возвращает новый набор данных вместо изменения существующего набора данных.