torch.distributions.Categorical

作用是创建以参数probs为标准的类别分布，样本是来自 “0 … K-1” 的整数，其中 K 是probs参数的长度。也就是说，按照传入的probs中给定的概率，在相应的位置处进行取样，取样返回的是该位置的整数索引。

如果 probs 是长度为 K 的一维列表，则每个元素是对该索引处的类进行抽样的相对概率。

如果 probs 是二维的，它被视为一批概率向量

如果是N维的，只有第N维被视为概率向量，前面维度被视为batch

torch.distributions.categorical.Categorical(probs=None, 
                                            logits=None, 
                                            validate_args=None)
创建一个离散的类别分布，参数由 probs 和 logits, 二者其一指定

示例一

probs = torch.FloatTensor([[0.05,0.1,0.85],[0,0.5,0.9]])
 
dist = Categorical(probs)
print(dist)
# Categorical(probs: torch.Size([2, 3]))
 
index = dist.sample()
print(index.numpy())
# [2 2]

示例二

probs = torch.FloatTensor([0.9,0.2])
D = Categorical(probs)
print(D)                  # Categorical(probs: torch.Size([2]))
for i in range(5):
    print(D.sample())         
tensor(0)
tensor(0)
tensor(0)
tensor(0)
tensor(1)
 
 
 
probs = torch.FloatTensor([[0, 0.25, 0.25], [0.25, 0.25, 0.5]])
D = Categorical(probs)
for i in range(5):
    print(D.sample())
 
tensor([1, 1])
tensor([2, 0])
tensor([2, 2])
tensor([2, 2])
tensor([2, 2])
# 第0维第0个元素的权重是0，在其他元素被取完之前是不会被取到的。

如果 probs 变量的 -1 维度上求和不为1，Categorical 内部也会帮你归一化，然后再 sample