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

Question

我正在设置一个 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() 输入，但这似乎绝对不会在多个线程中运行。任何建议都非常受欢迎。

Answer 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().

Answer 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 变量。