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

Posted 2021-05-07

我正在设置一个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）。另外，我已经查看了from_generator()的tf.data.Dataset输入，但这绝对不会在多线程中运行。任何建议都非常受欢迎。

答案

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

Generator definition

为了允许用户传入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块内的任何内容，以匹配数据集的构造方式和要获取的输出。

Usage example

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图像，因此您需要在上面的调用中修改这些以匹配您的用例。

Handling multiple files

我有一个处理多个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！

Possible caveats

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

在不同环境或分布式培训中应用此方法可能具有挑战性：

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

如果发电机依赖于外部状态，请小心：

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

另一答案

我花了一段时间来弄清楚这一点，所以我想我应该在这里记录下来。根据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变量。