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中采样一个值，其中值本身就是被采样的可能性。例如，从A采样0.0316的可能性是0.0316。采样值的输出应该仍然是一个张量。

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

让这件事变得棘手的一点是，我还想知道采样值在张量中的索引。也就是说，假设我采样0.2338，我想知道它是张量A的索引1、2还是3。

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})

Answer 2

这个怎么样？

probs = torch.tensor([0.0316, 0.2338, 0.2338, 0.2338, 0.0316, 0.0316, 0.0860, 0.0316, 0.0860],requires_grad=True)
dist=torch.distributions.categorical.Categorical(probs=probs)
probs[dist.sample()]

Answer 3

你可以通过这样做来作弊：

A = A*10000
temp = [[i]*A[i] for i in range(len(A))]
value = np.random.choice(temp)/10000