如何在tensorflow中填充序列

How to pad the sequence in tensorflow

我有一个包含像这样参差不齐的张量的数据集:

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]]>

有什么办法吗?谢谢大家!

您只需将必要的行数连接到每个参差不齐的张量即可:

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]]>