如何重复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

我们可以像下面这样做，

在使用flat_map.的 function.

1. Flatten
2. tf.repeat中，使用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]

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