PyTorch：如何从张量中采样，其中张量中的每个值都有不同的被选择可能性？

Question

给定 A = torch.tensor([0.0316, 0.2338, 0.2338, 0.2338, 0.0316, 0.0316, 0.0860, 0.0316, 0.0860])包含总和为 1 的概率的张量（我删除了一些小数，但可以安全地假设它总和为 1），我想采样一个值，其中A值本身就是采样的可能性。例如，0.0316从中采样的可能性A是0.0316。采样值的输出仍然应该是一个张量。

我尝试使用WeightedRandomSampler，但它不允许选择的值再成为张量，而是分离。

使这一点变得棘手的一个警告是，我还想知道张量中出现的采样值的索引。也就是说，假设我采样0.2338，我想知道它是索引1还是张2量。3A

Answer 1

可以通过累加权重并选择随机浮点[0,1)的插入索引来实现具有预期概率的选择。示例数组A稍微调整为总和为 1。

import torch

A = torch.tensor([0.0316, 0.2338, 0.2338, 0.2338, 0.0316, 0.0316, 0.0860, 0.0316, 0.0862], requires_grad=True)

p = A.cumsum(0)
#tensor([0.0316, 0.2654, 0.4992, 0.7330, 0.7646, 0.7962, 0.8822, 0.9138, 1.0000], grad_fn=<CumsumBackward0>))

idx = torch.searchsorted(p, torch.rand(1))
A[idx], idx

输出

(tensor([0.2338], grad_fn=<IndexBackward0>), tensor([3]))

这比更常见的方法更快A.multinomial(1)。
对一个元素采样 10000 次以检查分布是否符合概率

from collections import Counter

Counter(int(A.multinomial(1)) for _ in range(10000))
#1 loop, best of 5: 233 ms per loop

# vs @HatemAli's solution
dist=torch.distributions.categorical.Categorical(probs=A)
Counter(int(dist.sample()) for _ in range(10000))
# 10 loops, best of 5: 107 ms per loop

Counter(int(torch.searchsorted(p, torch.rand(1))) for _ in range(10000))
# 10 loops, best of 5: 53.2 ms per loop

输出

Counter({0: 319,
         1: 2360,
         2: 2321,
         3: 2319,
         4: 330,
         5: 299,
         6: 903,
         7: 298,
         8: 851})