torch.distributions.multinomial.
Multinomial
(total_count=1, probs=None, logits=None, validate_args=None)
probs是权重表,不是相加和为1也没关系,会自动帮你归一化,total_count就是要模拟多少次。sample后面可以加数值让他做多少批次。下面是掷骰子的例子
>>> fair_probs = torch.ones([6]) / 6
>>> multinomial.Multinomial(1, fair_probs).sample()
tensor([0., 0., 0., 1., 0., 0.])
>>> multinomial.Multinomial(10, fair_probs).sample((10,))
tensor([[3., 0., 2., 3., 0., 2.],
[0., 1., 6., 0., 1., 2.],
[4., 3., 1., 1., 1., 0.],
[2., 1., 0., 1., 1., 5.],
[2., 5., 1., 1., 0., 1.],
[0., 3., 1., 2., 4., 0.],
[1., 5., 0., 1., 2., 1.],
[4., 2., 1., 1., 1., 1.],
[1., 3., 2., 1., 3., 0.],
[0., 4., 3., 2., 1., 0.]])