python - 在 initializable_iterator 不可用的急切执行模式下，如何动态地提供 tf.data.Dataset？

Question

当我们需要逐个样本地提供数据时，以动态方式通过数据集管道提供数据的新方法(在急切执行下)是什么？

我有一个 tf.data.Dataset，它执行一些预处理步骤并从生成器读取数据，在训练期间从大型数据集中提取数据。

假设数据集表示为:

ds = tf.data.Dataset.from_tensor_slices([1, 2, 3, 4, 5, 6])
ds = ds.map(tf.square).shuffle(2).batch(2)
iterator = tf.data.make_one_shot_iterator(ds)

训练后，我想生成各种可视化效果，这需要我通过网络一次输入一个样本进行推理。我现在已经有了这个数据集预处理管道，我需要通过它来提供我的原始样本，以便为网络输入调整大小和形状。

这似乎是可初始化迭代器的一个用例:

placeholder = tf.placeholder(tf.float32, shape=None)
ds = tf.data.Dataset.from_tensor_slices(placeholder)
ds = ds.map(tf.square).shuffle(2).batch(2)
iterator = tf.data.make_initializable_iterator(ds) 
# now re-initialize for each sample

Keep in mind that the map operation in this example represents a long sequence of preprocessing operations that can't be duplicated for each new data sample being feed in.

这不适用于急切执行，您不能使用占位符。文档示例似乎都采用静态输入，例如此处的第一个示例。

我能想到的唯一方法是使用队列和 tf.data.Dataset.from_generator(...)，它在预测数据之前从我推送到的队列中读取.但这感觉既棘手，而且似乎容易出现我尚未解决的死锁。

TF 1.14.0

Answer 1

我刚刚意识到这个问题的答案很简单:

Just create a new dataset!

在 non-eager 模式下，下面的代码会降低性能，因为每个数据集操作都会被添加到图中并且永远不会释放，而在 non-eager 模式下，我们有可初始化的迭代器来解决这个问题。

然而，在急切执行模式下，像这样的 tensorflow 操作是短暂的，添加的迭代器不会被添加到全局图中，它们只是被创建并在不再被引用时消亡。为 TF2.0 赢得一个!

下面的代码(复制/粘贴可运行)演示:

import tensorflow as tf
import numpy as np
import time


tf.enable_eager_execution()

inp = np.ones(shape=5000, dtype=np.float32)

t = time.time()
while True:
    ds = tf.data.Dataset.from_tensors(inp).batch(1)
    val = next(iter(ds))
    assert np.all(np.squeeze(val, axis=0) == inp)
    print('Processing time {:.2f}'.format(time.time() - t))
    t = time.time()

提问的动机是在1.14的这个问题之后，在Keras下以图形模式创建多个数据集操作构成内存泄漏。

https://github.com/tensorflow/tensorflow/issues/30448