如何在tensorflow中填充序列

Question

我有一个包含粗糙张量的数据集，如下所示：

non_ragged_dataset = tf.data.Dataset.from_tensor_slices([1, 5, 3, 2, 8])
non_ragged_dataset = non_ragged_dataset.map(tf.range)
batched_non_ragged_dataset = non_ragged_dataset.apply(
    tf.data.experimental.dense_to_ragged_batch(2))
for element in batched_non_ragged_dataset:
    print(element)
----output----
<tf.RaggedTensor [[0], [0, 1, 2, 3, 4]]>
<tf.RaggedTensor [[0, 1, 2], [0, 1]]>
<tf.RaggedTensor [[0, 1, 2, 3, 4, 5, 6, 7]]>

每个RaggedTensor表示一个序列，形状为(None，None)。我想将序列转换为length=5，即shape=(5，None)，如下所示：

<tf.RaggedTensor [[0], [0, 1, 2, 3, 4], [0], [0], [0]]>
<tf.RaggedTensor [[0, 1, 2], [0, 1], [0], [0], [0]]>
<tf.RaggedTensor [[0, 1, 2, 3, 4, 5, 6, 7], [0], [0], [0], [0]]>

有什么办法吗？谢谢大家！

Answer 1

您可以将所需的行数连接到每个参差不齐的张量：

import tensorflow as tf

def pad_ragged_arrays(length):
    def transform(r):
        pad_size = tf.math.maximum(length - r.nrows(), 0)
        return tf.concat([r, tf.zeros([pad_size, 1], dtype=r.dtype)], axis=0)
    return transform

length = 5
dataset = (tf.data.Dataset.from_tensor_slices([1, 5, 3, 2, 8])
           .map(tf.range)
           .apply(tf.data.experimental.dense_to_ragged_batch(2))
           .map(pad_ragged_arrays(length)))
for element in dataset:
    print(element)
# <tf.RaggedTensor [[0], [0, 1, 2, 3, 4], [0], [0], [0]]>
# <tf.RaggedTensor [[0, 1, 2], [0, 1], [0], [0], [0]]>
# <tf.RaggedTensor [[0, 1, 2, 3, 4, 5, 6, 7], [0], [0], [0], [0]]>