如何重复 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 函数中使用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())

输出 -

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

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