Загрузка набора данных tf.data занимает много времени по сравнению с массивами numpy

Question

В настоящее время я занимаюсь поиском и устранением неисправностей / проверкой производительности моего конвейера ввода данных, и я очень удивлен, увидев, что загрузка данных в tf.data занимает много времени по сравнению с использованием numpy массивов. Некоторое время я играл с tf.data и пробовал несколько разных конфигураций, но для краткости я публикую только конвейер numpy и конвейер tf.data, которые (на мой взгляд) должны работать лучше всего. Рад опубликовать и другие конфигурации, хотя, если это поможет.
Кроме того, я сократил конвейеры до базовых c примеров без каких-либо дополнений и т. Д. c, просто читая изображения из файлов, нормализуя их, а затем создавая 2-кратное сокращение. версия этого. Уменьшенные версии - это данные обучения, а оригинальные изображения - метки обучения.
Я проверил, что с тяжелой предварительной обработкой / расширением или без нее, а также с обучением или без него результаты одинаковы. Конвейер numpy всегда загружает данные намного быстрее. Что касается времени эпохи во время фактического обучения, конвейеры tf.data также немного быстрее, но я думаю, что я должен сделать отдельный пост для этого.

Вот различные конвейеры

class NumpyPipeline:
    def __init__(self, hr_img_path):
        self.train_list = os.listdir(hr_img_path)
        imgs = [NumpyPipeline._read_hr_img(f"{hr_img_path}/{fname}") for fname in self.train_list]
        num_imgs = len(imgs)

        # very clunky, but a lot faster compared to converting/reshaping lists to np.arrays
        lr_imgs = np.zeros(shape=(num_imgs, 64, 64, 3), dtype=np.float32)
        idx = 0
        for img in imgs:
            lr_imgs[idx] = NumpyPipeline._prepare_img_pairs(img, (64, 64))
            idx += 1

        self.hr_imgs = imgs
        self.lr_imgs = lr_imgs

    def batch_generator(self, batch_size=8, seed=None):
        i = 0
        shuffled_lr_imgs, shuffled_hr_imgs = sklearn.utils.shuffle(
            self.lr_imgs, self.hr_imgs, random_state=seed
        )
        num_samples = len(self.lr_imgs)
        while i < num_samples:
            yield shuffled_lr_imgs[i:i+batch_size], shuffled_hr_imgs[i:i+batch_size]
            i += batch_size

    @staticmethod
    def init_pipeline(data_path, batch_size=None):
        return NumpyPipeline(hr_img_path=data_path)

    @staticmethod
    def _read_hr_img(fpath):
        img = tf.io.read_file(fpath)
        img = tf.image.decode_png(img)
        img = tf.image.convert_image_dtype(img, tf.float32)
        return img

    @staticmethod
    def _prepare_img_pairs(hr_img, lr_dims):
        return tf.image.resize(hr_img, size=lr_dims) 


class VectorizedMappingPipeline:
    def __init__(self, hr_img_path, batch_size=8):
        self.train_list = os.listdir(hr_img_path)
        train_list_ds = tf.data.Dataset.list_files([f"{hr_img_path}/{fname}" for fname in self.train_list])

        # shuffle early
        train_list_ds = train_list_ds.shuffle(4096)

        # read files in parallel
        train_hr_ds = train_list_ds.map(
            VectorizedMappingPipeline._read_hr_img,
            num_parallel_calls=tf.data.experimental.AUTOTUNE
        )

        # batch early for vectorized mapping
        train_hr_ds = train_hr_ds.batch(batch_size)
        train_hr_ds = train_hr_ds.map(
            lambda hr_img: VectorizedMappingPipeline._prepare_img_pairs(hr_img, (64, 64)),
            num_parallel_calls=tf.data.experimental.AUTOTUNE
        )

        self.train_ds = train_hr_ds

    def batch_generator(self, batch_size=None):
        return self.train_ds.prefetch(tf.data.experimental.AUTOTUNE)

    @staticmethod
    def init_pipeline(data_path, batch_size=8):
        return VectorizedMappingPipeline(hr_img_path=data_path, batch_size=batch_size)

    @staticmethod
    def _read_hr_img(fpath):
        img = tf.io.read_file(fpath)
        img = tf.image.decode_png(img)
        img = tf.image.convert_image_dtype(img, tf.float32)
        return img

    @staticmethod
    def _prepare_img_pairs(hr_img, lr_dims):
        lr_img = tf.image.resize(hr_img, size=lr_dims)
        return lr_img, hr_img

Моя функция бенчмарка очень проста и выглядит следующим образом:

    start_loading_pipeline = time.perf_counter()
    data_pipeline = data_pipeline.init_pipeline(DATA_DIR, BATCH_SIZE)
    stop_loading_pipeline = time.perf_counter()
    return stop_loading_pipeline - start_loading_pipeline

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

Вот результаты для разных размеров набора данных (я также пробовал использовать комбинации больших размеров наборов данных / размеров пакетов, результаты всегда одинаковы):

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