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

、

3

或

A

。

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

上面接受的答案的解决方案可以扩展到覆盖具有概率的二维张量，例如模型输出的 softmax。只需相应地调整随机张量的尺寸即可。

def multisample_from_softmax(softmax_values):
    """
    quick weighted sampling using pytorch
    softmax_values : torch.tensor shaped (n_tokens, embedding_vocab_size)
    returns: torch.tensor shaped(n_tokens) with indices of sampled tokens
    """
    size = softmax_values.shape[0]
    rand_values = torch.rand((size, 1), device=softmax_values.device)
    cumprobs = softmax_values.cumsum(dim=1)
    selection = torch.searchsorted(cumprobs, rand_values).squeeze(1)
    selection_probs = (softmax_values[:, selection] * torch.eye(size, device=softmax_values.device)).diagonal()
    return selection, selection_probs

Answer 4

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

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