python - 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 block 中的任何内容，以匹配您的数据集的构建方式以及您想要获得的输出。

使用示例

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

同样，在我的例子中，我存储了 uint8 高度 427、宽度 561 和 3 的图像我的数据集中的颜色 channel ，因此您需要在上述调用中修改它们以匹配您的用例。

处理多个文件

我有一个处理多个 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提到这两个注释。

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

NOTE: The current implementation of Dataset.from_generator() uses tf.py_func and inherits the same constraints. In particular, it requires the Dataset- and Iterator-related operations to be placed on a device in the same process as the Python program that called Dataset.from_generator(). The body of generator will not be serialized in a GraphDef, and you should not use this method if you need to serialize your model and restore it in a different environment.

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

NOTE: If generator depends on mutable global variables or other external state, be aware that the runtime may invoke generator multiple times (in order to support repeating the Dataset) and at any time between the call to Dataset.from_generator() and the production of the first element from the generator. Mutating global variables or external state can cause undefined behavior, and we recommend that you explicitly cache any external state in generator before calling Dataset.from_generator().