torch.multinomial与torch.distributions.categorical.Categorical功能是否一致？

首先给出明确结论：二者采样逻辑完全等价，当你传入[0.5, 0.5]的概率数组时，都会以均等概率采样得到0和1的索引。

核心关联 #

• torch.distributions.categorical.Categorical的采样方法底层直接调用torch.multinomial实现，二者采样的分布输出没有任何统计差异
• 两者都支持传入未归一化的概率值或者对数概率值，内部会自动完成归一化处理，不需要手动保证概率和为1

代码验证示例 #

你可以运行以下代码验证二者的采样结果一致性：

import torch

# 定义目标概率分布
probs = torch.tensor([0.5, 0.5])
sample_num = 100000

# 用torch.multinomial采样
multinomial_res = torch.multinomial(probs, num_samples=sample_num, replacement=True)
multinomial_0_ratio = (multinomial_res == 0).float().mean().item()
print(f"multinomial采样0的占比: {multinomial_0_ratio:.4f}")

# 用Categorical采样
cat_dist = torch.distributions.Categorical(probs=probs)
cat_res = cat_dist.sample((sample_num,))
cat_0_ratio = (cat_res == 0).float().mean().item()
print(f"Categorical采样0的占比: {cat_0_ratio:.4f}")

运行后两个采样结果中0的占比都会在0.5附近小范围波动，属于正常的统计误差范围。

二者的使用差异 #

二者只是封装层级不同，适用场景有区别：

• torch.multinomial是更底层的API，仅提供采样功能，适合只需要获取采样索引的简单场景
• Categorical是概率分布封装类，除采样外还支持直接计算对数似然、分布熵、不同分布间的KL散度等操作，更适合强化学习、生成模型等需要频繁操作分布属性的场景

注意：如果传入的是对数概率，两个API都可以指定logits参数传入，避免手动计算softmax带来的数值不稳定问题。

内容的提问来源于stack exchange，提问作者Mzq