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