与线程/队列相比，tf.data.Dataset 输入管道提供了糟糕的结果

Question

之前我使用线程和队列作为我的数据管道，我在两个 GPU 上都获得了非常高的 Util（数据是动态创建的）。我想使用 tf 数据集，但我很难重现结果。

我尝试了很多方法。因为我动态创建数据，所以 from_generator() 方法看起来很完美。您在下面看到的代码是我最后一次尝试。尽管我使用 map() 函数来处理生成的图像，但似乎在创建数据方面存在瓶颈。我在下面的代码中尝试了什么，我想以某种方式“多线程”生成器，因此同时有更多数据进入。但目前还没有更好的结果。

def generator(n):
    with tf.device('/cpu:0'):
        while True:
            ...
            yield image, label

def get_generator(n):
    return partial(generator, n)

def dataset(n):
    return tf.data.Dataset.from_generator(get_generator(n), output_types=(tf.float32, tf.float32), output_shapes=(tf.TensorShape([None,None,1]),tf.TensorShape([None,None,1])))

def input_fn():
# ds = tf.data.Dataset.from_generator(generator, output_types=(tf.float32, tf.float32), output_shapes=(tf.TensorShape([None,None,1]),tf.TensorShape([None,None,1])))
    ds = tf.data.Dataset.range(BATCH_SIZE).apply(tf.data.experimental.parallel_interleave(dataset, cycle_length=BATCH_SIZE))
    ds = ds.map(map_func=lambda img, lbl: processImage(img, lbl))
    ds = ds.shuffle(SHUFFLE_SIZE)
    ds = ds.batch(BATCH_SIZE)
    ds = ds.prefetch(1)
return ds

预期结果将是较高的 GPU Util (>80%)，但目前确实很低 10/20%。

Answer 1

您可以改用tf.data.Dataset.from_tensor_slices。 只需传递图像/标签路径。该函数接受文件名作为参数。

def input_func():
    dataset = tf.data.Dataset.from_tensor_slices(images_path, labels_path)
    dataset = dataset.shuffle().repeat()
    ...
    return dataset