如何将具有不同文件数量的多个目录读入tensorflow数据集中

Question

我的 tensorflow 版本是 1.15，我的目录树如下图所示。

root
|
|---scene1
|     |
|     |--img1.npy
|     |--img2.npy
|     |--cam.txt
|     |--poses.txt
|
|---scene2
|     |
|     |--img1.npy
|     |--img2.npy
|     |--img3.npy
|     |--cam.txt
|     |--poses.txt

每个场景文件夹包含不同数量的图像（npy 格式），但只有一个 cam.txt 和一个poses.txt。我尝试过使用numpy.genfromtxt 和numpy.load 将每个场景文件夹中的文件读入张量，然后使用

ds = tf.data.Dataset.from_tensors 为每个场景创建数据集，最后使用ds.concatenate 连接这些数据集。这种方法有效，但是当场景文件夹的数量很大时会浪费很多时间。有没有更好的方法来处理这个问题？

Answer 1

我最近遇到了类似的问题。

在处理非常大的数据集时，Python generators 是一个不错的选择：

[Generators] 像常规函数一样编写，但只要它们想要返回数据就使用 yield 语句。每次在其上调用 next() 时，生成器都会从中断处继续（它会记住所有数据值以及上次执行的语句）。

Tensorflow 的 Dataset 类通过静态 from_generator 方法支持它们（也可以使用 with TF 1.15）：

import pathlib
import numpy as np
import tensorflow as tf

def scene_generator():
    base_dir = pathlib.Path('path/to/root/')
    scenes = tf.data.Dataset.list_files(str(base_dir / '*'))
    for s in scenes:
        scene_dir = pathlib.Path(s.numpy().decode('utf-8'))
        images = scene_dir / '*.npy'
        data = []
        for i in images:
            with np.load(i) as image:
                features = [image['x1'], image['x2']] # ...
                data.append(features)
        cam   = np.genfromtxt(scene_dir / 'cam.txt')
        poses = np.genfromtxt(scene_dir / 'poses.txt')

        yield data, [cam, poses]

types = (tf.float32, tf.float32)
shapes = (tf.TensorShape([None]), tf.TensorShape([2, None]))
train_data = tf.data.Dataset.from_generator(scene_generator,
    output_types=types,
    output_shapes=shapes)