何时在 Tensorflow Estimator 中使用迭代器

Question

在 Tensorflow 指南中，指南有两个不同的地方描述了 Iris Data 示例的输入函数。一个输入函数只返回数据集本身，而另一个返回带有迭代器的数据集。

来自预制的 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_resulthere。

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