来自粗糙张量的样本

Question

我有一个从1到10k的row_lens的raggedTensor。我想以可伸缩的方式从中随机选择元素，每行的数量有上限。如本例所示：

vect = [[1,2,3],[4,5][6],[7,8,9,10,11,12,13]]
limit = 3
sample(vect, limit)

->输出：[[1,2,3]，[4,5]，[6]，[7,9,11]]

我的想法是在len_row<在另一种情况下是随机的。我想知道这是否可以通过一些tensorflow操作以低于batch_size的复杂性来完成？

Answer 1

您可以尝试在图形模式下使用tf.map_fn：

import tensorflow as tf
vect = tf.ragged.constant([[1,2,3],[4,5],[6],[7,8,9,10,11,12,13]])
@tf.function
def sample(x, samples=3):
length = tf.shape(x)[0]
x = tf.cond(tf.less_equal(length, samples), lambda: x, lambda: tf.gather(x, tf.random.shuffle(tf.range(length))[:samples]))
return x
c = tf.map_fn(sample, vect)

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

请注意，tf.vectorized_map可能会更快，但目前存在关于该函数和粗糙张量的错误。使用tf.while_loop也是一种选择。