tensorflow - 何时在 Tensorflow Estimator 中使用迭代器

Question

在 Tensorflow 指南中有两个单独的地方，指南描述了虹膜数据示例的输入函数。一个输入函数仅返回数据集本身，而另一个输入函数返回带有迭代器的数据集。

来自预制的 Estimator 指南:https://www.tensorflow.org/guide/premade_estimators

def train_input_fn(features, labels, batch_size):
"""An input function for training"""
# Convert the inputs to a Dataset.
dataset = tf.data.Dataset.from_tensor_slices((dict(features), labels))

# Shuffle, repeat, and batch the examples.
return dataset.shuffle(1000).repeat().batch(batch_size)

来自自定义估算器指南: https://www.tensorflow.org/guide/custom_estimators

def train_input_fn(features, labels, batch_size):
"""An input function for training"""
# Convert the inputs to a Dataset.
dataset = tf.data.Dataset.from_tensor_slices((dict(features), labels))

# Shuffle, repeat, and batch the examples.
dataset = dataset.shuffle(1000).repeat().batch(batch_size)

# Return the read end of the pipeline.
return dataset.make_one_shot_iterator().get_next()

我很困惑哪个是正确的，如果它们都用于不同的情况，什么时候使用迭代器返回数据集是正确的？

Answer 1

如果您的输入函数返回 tf.data.Dataset ，在引擎盖下创建了一个迭代器及其 get_next()函数用于为模型提供输入。这在源代码中有些隐藏，见parse_input_fn_result here .

我相信这仅在最近的更新中实现，因此较旧的教程仍然明确返回 get_next()在他们的输入功能中，因为它是当时唯一的选择。使用两者应该没有区别，但是您可以通过返回数据集而不是迭代器来节省一点代码。