PyTorch中nn.softmax的dim=0/1参数使用规则疑问

结论先行 #

你对dim=0沿列计算、dim=1沿行计算的判断完全正确。

PyTorch dim参数通用逻辑 #

PyTorch中所有需要指定dim参数的运算，核心规则都是：运算会沿着你指定的dim维度，在该维度的所有元素之间进行计算，计算后该维度的求和结果为1（对应softmax这类归一化运算）。
对于二维张量来说，shape格式为(行数, 列数)：

• dim=0对应第一个维度，也就是「行维度」：运算时会遍历所有行，固定其他维度的索引不变，对同一位置的行元素做计算，落到你给出的示例里，就是每一列的所有行元素做softmax，结果每列和为1
• dim=1对应第二个维度，也就是「列维度」：运算时会遍历所有列，固定其他维度的索引不变，对同一位置的列元素做计算，落到你给出的示例里，就是每一行的所有列元素做softmax，结果每行和为1

示例验证 #

你给出的测试代码完全符合上述逻辑：

>>> import torch
>>> import torch.nn.functional as F
>>> x = torch.tensor([[1,2],[3,4]],dtype=torch.float)
# dim=0沿行维度计算，每列和为1
>>> F.softmax(x,dim=0)
tensor([[0.1192, 0.1192],
        [0.8808, 0.8808]])
# dim=1沿列维度计算，每行和为1
>>> F.softmax(x,dim=1)
tensor([[0.2689, 0.7311],
        [0.2689, 0.7311]])

可以手动验算验证：

• dim=0时第一列元素为1、3，softmax计算结果为e^1/(e^1+e^3) ≈ 0.1192、e^3/(e^1+e^3)≈0.8808，和输出一致
• dim=1时第一行元素为1、2，softmax计算结果为e^1/(e^1+e^2)≈0.2689、e^2/(e^1+e^2)≈0.7311，和输出一致

通用判断技巧 #

不管是几维张量，你只需要记住：你想要哪个维度的归一化和为1，就把dim设置为对应维度的序号即可。最常见的分类场景中，模型输出shape为(batch_size, num_classes)，我们需要每个样本对应所有类别的概率和为1，因此默认设置dim=1。

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