我应该直接返回数据集还是应该使用 one_shot 迭代器？

【问题标题】：Should I return dataset directly or should i use one_shot iterator instead?我应该直接返回数据集还是应该使用 one_shot 迭代器？
【发布时间】：2019-02-04 16:11:45
【问题描述】：

我正在使用 Dataset API 构建数据管道，但是当我训练多个 GPU 并在我的输入函数中返回 dataset.make_one_shot_iterator().get_next() 时，我得到了

ValueError: dataset_fn() must return a tf.data.Dataset when using a tf.distribute.Strategy

我可以按照错误消息直接返回数据集，但我不明白 iterator().get_next() 的用途以及它如何在单 GPU 和多 GPU 上进行训练。

...

    dataset = dataset.repeat(num_epochs)
    dataset = dataset.batch(batch_size = batch_size)
    dataset = dataset.cache()

    dataset = dataset.prefetch(buffer_size=None)

    return dataset.make_one_shot_iterator().get_next()

return _input_fn

【问题讨论】：

标签： python tensorflow iterator pipeline tensorflow-datasets

【解决方案1】：

将tf.data 与分发策略一起使用时（可与Keras 和tf.Estimators 一起使用），您的输入fn 应返回tf.data.Dataset：

def input_fn():
  dataset = dataset.repeat(num_epochs)
  dataset = dataset.batch(batch_size = batch_size)
  dataset = dataset.cache()

  dataset = dataset.prefetch(buffer_size=None)
  return dataset

...use input_fn...

请参阅documentation 了解分发策略。

dataset.make_one_shot_iterator() 在分发策略/更高级别的库之外很有用，例如，如果您正在使用更低级别的库，或者调试/测试数据集。例如，您可以像这样迭代数据集的所有元素：

dataset = ...
iterator = dataset.make_one_shot_iterator()
get_next = iterator.get_next()
with tf.Session() as sess:
  while True:
    print(sess.run(get_next))
  except tf.errors.OutOfRangeError:
    break

【讨论】：