【torch函数】torch.multinomial函数

torch.multinomial

是PyTorch中的一个函数，用于从多项分布中抽取样本。多项分布是一种描述多个可能结果的概率分布，例如抛硬币、掷骰子等。

torch.multinomial的用法如下：

torch.multinomial(input, num_samples, replacement=False, *, generator=None, out=None)

参数说明：

input: 输入张量，表示每个类别的概率分布。它可以是一个一维张量，也可以是一个二维张量。如果是一维张量，它表示每个类别的概率；如果是二维张量，每行表示一个概率分布。

num_samples: 抽样次数，即要从每个概率分布中抽取的样本数量。

replacement: 是否进行有放回抽样。默认为False，即不进行有放回抽样。

generator: 随机数生成器。如果为None，将使用默认的随机数生成器。

out: 输出张量，用于存储抽取的样本。

函数返回一个包含抽样结果的张量，其形状为(input.size(0), num_samples)。每行表示从对应概率分布中抽取的样本。

下面是一个简单的示例：

import torch

# 定义概率分布
probs = torch.tensor([0.1, 0.2, 0.3, 0.4])

# 从概率分布中抽取样本
samples = torch.multinomial(probs, num_samples=5, replacement=True)

print(samples)

此示例从给定的概率分布中抽取了5个样本，其中每个样本是从0到3的整数，表示概率分布中对应类别的索引。