使用 TensorFlow Data API 和 Keras 模型预测批次

Question

假设我有一个数据集和一个 Keras 模型。数据集已使用 tf Dataset API 中的batch() 分批。现在我正在寻找一种高效且干净的方法来对所有测试样本进行批量预测。

我已经尝试了以下代码并且它有效。

batch_size = 32
dataset = dataset.batch(batch_size)
predictions = keras_model.predict(dataset, steps=math.ceil(num_testing_samples / batch_size))

我想知道有没有更有效和优雅的方法来实现这一点？

Answer 1

TF >= 1.14.0

你可以设置steps=None。来自tf.keras.Model.predict()的官方文档：

如果 x 是 tf.data 数据集且步骤为 None，则 predict 将运行直到输入数据集用完。

只要确保您的 dataset 对象未处于重复模式，您就可以开始了 :)。

TF 1.12.0 & 1.13.0

这些版本对tf.data.Dataset 和tf.keras 的支持非常差。 tf.data.Dataset 对象被转换为一个迭代器 here，如果你没有设置 steps 参数，它会触发一个错误 here。这是在 1.14.0 中修补的。