在张量流中，沿 0 轴从张量中随机(子)采样 k 个条目

Question

给定rank>=1T的张量，我想沿 0 轴均匀地从中随机抽取k条目。

编辑：采样应该是计算图的一部分，作为惰性操作，并且每次调用时都应输出不同的随机条目。

例如，给定等级2的T：

T = tf.constant( 
[[1,1,1,1,1],
[2,2,2,2,2],
[3,3,3,3,3],
....
[99,99,99,99,99],
[100,100,100,100,100]] 
)

对于k=3，可能的输出为：

#output = 
#   [[34,34,34,34,34],
#    [6,6,6,6,6],
#    [72,72,72,72,72]]

如何在张量流中实现这一点？

Answer 1

您可以在索引数组上使用随机洗牌：

获取前sample_num索引，并使用它们来选择输入的切片。

idxs = tf.range(tf.shape(inputs)[0])
ridxs = tf.random.shuffle(idxs)[:sample_num]
rinput = tf.gather(inputs, ridxs)