如何重复 tf.data.Dataset

Question

tf.data.Dataset实际上有一个repeat方法，它输出更像一个tile的东西，即：

list(tf.data.Dataset.range(2).repeat(3).as_numpy_iterator())
# [0, 1, 0, 1, 0, 1]

就好像：

np.tile(np.arange(2), 3)
# array([0, 1, 0, 1, 0, 1])

我正在寻找的是实际的 numpy 重复：

np.repeat(np.arange(2), 3)
# array([0, 0, 0, 1, 1, 1])

Answer 1

我们可以这样做，如下所示，

1. 在map function 中使用tf.repeat重复每个元素。
2. 使用flat_map平结果。

代码 -

%tensorflow_version 2.x
import tensorflow as tf

dataset = tf.data.Dataset.range(2).map(lambda x : tf.repeat(x,3)).flat_map(lambda y: tf.data.Dataset.from_tensor_slices(y))

list(dataset.as_numpy_iterator())

Output -

[0, 0, 0, 1, 1, 1]

希望这能回答你的问题。 快乐学习。