如何用TensorFlow从概率分布张量中采样指定数量的索引及其概率

给定每行和为1的Softmax概率分布张量，要完成采样k个索引并获取对应概率值的需求，可以通过以下TensorFlow操作实现：

实现步骤 #

1. 采样索引：使用tf.random.categorical函数，该函数接受对数概率作为输入，因此需要先对输入的概率张量取自然对数。指定采样数量num_samples=3，即可得到每行3个采样索引。
2. 提取对应概率：使用tf.batch_gather函数，根据采样得到的索引，从原概率张量中批量提取对应位置的概率值。

完整代码示例 #

import tensorflow as tf

# 定义输入概率张量
probs = tf.constant([
    [0.0, 0.1, 0.2, 0.3, 0.4],
    [0.5, 0.3, 0.2, 0.0, 0.0]])

# 采样3个索引（每行）
sampled_indices = tf.random.categorical(tf.math.log(probs), num_samples=3)
# 根据索引提取对应概率
sampled_probs = tf.batch_gather(probs, sampled_indices)

# 打印结果
print("index:", sampled_indices.numpy())
print("\nprobs:", sampled_probs.numpy())

说明 #

• tf.random.categorical会按照每行的概率分布独立进行采样，因此每次运行的结果会略有不同，但符合给定的概率分布（比如第一行中索引4的概率最高，被采样到的次数会最多）。
• tf.batch_gather专门用于批量索引提取，能够高效地从二维张量的每行中按索引取值，完美匹配需求。

内容的提问来源于stack exchange，提问作者n0obcoder