TensorFlow - tf.data.Dataset 读取大型 HDF5 文件

Posted 2023-02-23

【中文标题】TensorFlow - tf.data.Dataset 读取大型 HDF5 文件

【英文标题】：TensorFlow - tf.data.Dataset reading large HDF5 files

【发布时间】：2018-06-26 20:05:26

【问题描述】：

我正在设置一个 TensorFlow 管道，用于读取大型 HDF5 文件作为我的深度学习模型的输入。每个 HDF5 文件包含 100 个可变大小长度的视频，这些视频存储为压缩 JPG 图像的集合（以使磁盘上的大小易于管理）。使用tf.data.Dataset 和tf.py_func 的映射，使用自定义Python 逻辑从HDF5 文件中读取示例非常容易。例如：

def read_examples_hdf5(filename, label):
    with h5py.File(filename, 'r') as hf:
        # read frames from HDF5 and decode them from JPG
    return frames, label

filenames = glob.glob(os.path.join(hdf5_data_path, "*.h5"))
labels = [0]*len(filenames) # ... can we do this more elegantly?

dataset = tf.data.Dataset.from_tensor_slices((filenames, labels))
dataset = dataset.map(
    lambda filename, label: tuple(tf.py_func(
        read_examples_hdf5, [filename, label], [tf.uint8, tf.int64]))
)

dataset = dataset.shuffle(1000 + 3 * BATCH_SIZE)
dataset = dataset.batch(BATCH_SIZE)
iterator = dataset.make_one_shot_iterator()
next_batch = iterator.get_next()

此示例有效，但问题是tf.py_func 似乎一次只能处理一个示例。由于我的 HDF5 容器存储了 100 个示例，因此此限制会导致大量开销，因为文件需要不断地打开、读取、关闭和重新打开。将所有 100 个视频示例读入数据集对象，然后继续处理下一个 HDF5 文件（最好在多个线程中，每个线程处理自己的 HDF5 文件集合）会更有效。

所以，我想要的是在后台运行多个线程，从 HDF5 文件中读取视频帧，从 JPG 解码它们，然后将它们输入数据集对象。在引入tf.data.Dataset 管道之前，使用RandomShuffleQueue 和enqueue_many 操作非常容易，但目前似乎没有优雅的方法来做到这一点（或者缺少文档）。

有谁知道实现我的目标的最佳方式是什么？我还使用tfrecord 文件研究（并实现了）管道，但是随机抽取存储在tfrecord 文件中的视频帧样本似乎是不可能的（请参阅here）。此外，我查看了tf.data.Dataset 的from_generator() 输入，但这似乎绝对不会在多个线程中运行。任何建议都非常受欢迎。

【问题讨论】：

tf.data.Dataset.map(your_map_function, num_parallel_calls=N) 做你想做的事吗？它将运行您的地图函数的N 线程。我看到的问题是您现在有 6 个线程，每个线程读取 1 个 HDF5 文件，这意味着您最好有足够的内存来存储所有 6 个完整的 HDF5 文件。我之所以提出这个问题，是因为我发布了一个相关问题，试图解决内存有限和大 HDF5 文件的问题。 ***.com/questions/48349409/…

HDF5 文件的内存不是问题，因为它不会将整个文件读入内存。但是始终打开-读取-关闭-重新打开...文件仍然存在问题，这将大大降低速度。

【参考方案1】：

我在处理类似问题时偶然发现了这个问题。我想出了一个基于使用 Python 生成器的解决方案，以及 TF 数据集构造方法from_generator。因为我们使用了生成器，所以 HDF5 文件应该只打开一次以供读取，并且只要有要读取的条目就保持打开状态。因此，它不会在每次调用时都打开、读取然后关闭以获取下一个数据元素。

发电机定义

为了允许用户将 HDF5 文件名作为参数传递，我生成了一个具有 __call__ 方法的类，因为 from_generator 指定生成器必须是可调用的。这是生成器：

import h5py
import tensorflow as tf

class generator:
    def __init__(self, file):
        self.file = file

    def __call__(self):
        with h5py.File(self.file, 'r') as hf:
            for im in hf["train_img"]:
                yield im

通过使用生成器，代码应该从上次返回结果时每次调用中断的地方重新开始，而不是从头开始运行所有内容。在这种情况下，它在内部 for 循环的下一次迭代中。所以这应该跳过再次打开文件进行读取，只要有数据到yield，就保持打开状态。有关生成器的更多信息，请参阅this excellent Q&A。

当然，您必须替换 with 块内的任何内容，以匹配您的数据集的构造方式以及您想要获得的输出。

使用示例

ds = tf.data.Dataset.from_generator(
    generator(hdf5_path), 
    tf.uint8, 
    tf.TensorShape([427,561,3]))

value = ds.make_one_shot_iterator().get_next()

# Example on how to read elements
while True:
    try:
        data = sess.run(value)
        print(data.shape)
    except tf.errors.OutOfRangeError:
        print('done.')
        break

再次，在我的例子中，我在我的数据集中存储了高度为427、宽度为561 和3 颜色通道的uint8 图像，因此您需要在上述调用中修改这些以匹配您的用例。

处理多个文件

我有一个处理多个 HDF5 文件的建议解决方案。基本思想是照常从文件名中构造一个Dataset，然后使用interleave 方法同时处理许多输入文件，例如从每个文件中获取样本以形成一个批处理。

思路如下：

ds = tf.data.Dataset.from_tensor_slices(filenames)
# You might want to shuffle() the filenames here depending on the application
ds = ds.interleave(lambda filename: tf.data.Dataset.from_generator(
        generator(filename), 
        tf.uint8, 
        tf.TensorShape([427,561,3])),
       cycle_length, block_length)

它的作用是同时打开 cycle_length 文件，并在移动到下一个文件之前从每个文件中生成 block_length 项目 - 有关详细信息，请参阅 interleave 文档。您可以在此处设置值以匹配适合您的应用程序的值：例如，您需要一次处理一个文件还是同时处理多个文件，您是否只想一次从每个文件中获取一个样本，等等.

编辑：对于并行版本，请查看tf.contrib.data.parallel_interleave！

可能的警告

如果您决定采用该解决方案，请注意使用 from_generator 的特殊性。对于 Tensorflow 1.6.0，documentation of from_generator 提到了这两个注释。

在不同的环境或分布式训练中应用它可能具有挑战性：

注意：Dataset.from_generator() 的当前实现使用 tf.py_func 并继承了相同的约束。特别是，它 需要放置与数据集和迭代器相关的操作 与调用的 Python 程序在同一进程中的设备 数据集.from_generator()。生成器的主体不会被序列化 在 GraphDef 中，如果需要，则不应使用此方法 序列化您的模型并在不同的环境中恢复它。

如果生成器依赖于外部状态，请注意：

注意：如果生成器依赖于可变的全局变量或其他 外部状态，请注意运行时可能会调用生成器 多次（为了支持重复数据集）并且在任何时候 调用 Dataset.from_generator() 和生产之间的时间 来自生成器的第一个元素。改变全局变量或 外部状态可能导致未定义的行为，我们建议您 在调用之前显式缓存生成器中的任何外部状态 Dataset.from_generator().

【讨论】：

有没有办法从同一个 hdf 文件中加载多个数据集？例如 hf["train_img"] 和 hf["labels"]。

我对此有一个后续问题。单个 HDF5 文件的示例可以正常工作，但我无法让多个示例与交错一起使用。问题是 tf.data.Dataset.from_tensor_slices(filenames) 返回张量对象的集合而不是 Python 字符串，因此生成器无法处理这个问题。处理这个问题的正确方法是什么？

也许您想将此添加到您的警告中（从 TF 2.6.x 开始）：“警告：虽然这是一种方便的方法，但它的可移植性和可扩展性有限。它必须在同一个 python 进程中运行它创建了生成器，并且仍然受制于 Python GIL。” tensorflow.org/guide/data#consuming_python_generators

【参考方案2】：

我花了一段时间才弄清楚这一点，所以我想我应该在这里记录一下。根据 mikkola 的回答，这是处理多个文件的方法：

import h5py
import tensorflow as tf

class generator:
    def __call__(self, file):
        with h5py.File(file, 'r') as hf:
            for im in hf["train_img"]:
                yield im

ds = tf.data.Dataset.from_tensor_slices(filenames)
ds = ds.interleave(lambda filename: tf.data.Dataset.from_generator(
        generator(), 
        tf.uint8, 
        tf.TensorShape([427,561,3]),
        args=(filename,)),
       cycle_length, block_length)

关键是您不能将filename 直接传递给generator，因为它是Tensor。你必须通过 args 传递它，它会被 tensorflow 评估并将其转换为常规的 Python 变量。

【讨论】：

我可以确认我需要这个修改才能让它工作。

您好，在上面的生成器中，您没有返回图片的标签吗？如何退回标签？我想要做的是：yield h5file[data_name]，标签，是否需要对返回的形状进行任何更改：tf.TensorShape([427,561,3]),

我很遗憾没有看第二个答案（即这个）。花了我几个小时来尝试获得原始作品！

@void 获取标签，我正在做的是将标签与生成器中的其余输出连接起来，然后通过 tf.data.Dataset 传递它。 map() 将其分成两部分，特征和标签。