如何在TensorFlow 2.0中使用padded_batch()

Question

X = tf.range(10)
dataset = tf.data.Dataset.from_tensor_slices(X)
dataset2 = dataset.repeat(3).padded_batch(7, padded_shapes=([]))
for item in dataset2:
    print(item)

输出

tf.Tensor([0 1 2 3 4 5 6], shape=(7,), dtype=int32)
tf.Tensor([7 8 9 0 1 2 3], shape=(7,), dtype=int32)
tf.Tensor([4 5 6 7 8 9 0], shape=(7,), dtype=int32)
tf.Tensor([1 2 3 4 5 6 7], shape=(7,), dtype=int32)
tf.Tensor([8 9], shape=(2,), dtype=int32)

如何定义padded_shapes来获得像下面这样的结果？

tf.Tensor([0 1 2 3 4 5 6], shape=(7,), dtype=int32)
tf.Tensor([7 8 9 0 1 2 3], shape=(7,), dtype=int32)
tf.Tensor([4 5 6 7 8 9 0], shape=(7,), dtype=int32)
tf.Tensor([1 2 3 4 5 6 7], shape=(7,), dtype=int32)
tf.Tensor([8 9 0 0 0 0 0], shape=(7,), dtype=int32)

Answer 1

我用batch(7)解决了这个问题。

dataset2 = dataset.repeat(3).batch(7).padded_batch(7, padded_shapes=([None]))

输出

tf.Tensor(
[[0 1 2 3 4 5 6]
 [7 8 9 0 1 2 3]
 [4 5 6 7 8 9 0]
 [1 2 3 4 5 6 7]
 [8 9 0 0 0 0 0]], shape=(5, 7), dtype=int32)