为何在torch.softmax()中使用dim=0获取预测概率？

先明确PyTorch里dim参数的核心逻辑：dim指定的是你要在哪个维度上执行计算（聚合/归一化），不同的张量维度定义，对应的dim选择会不一样。

为什么教程用dim=0？你该用dim=1吗？ #

这取决于你的张量维度定义：

• 如果你的输入logits形状是(N, C)（N=样本数，C=类别数），每一行对应一个样本的所有类别logits，那确实应该用torch.softmax(dim=1)——对每个样本的类别logits做归一化，让每行的概率和为1，这也是分类任务的常规用法，和你说的torch.argmax(..., dim=1)完全匹配（argmax找每行里最大概率的类别索引）。
• 但如果教程里的logits形状是(C, N)（类别数在前，样本数在后），每一列对应一个样本的所有类别logits，那用dim=0才是对的——沿着类别维度（第0维）做softmax，让每列的概率和为1，和torch.argmax(..., dim=0)对应（找每列里最大的类别索引）。

简单说：类别维度在哪个位置，就对哪个dim做softmax/argmax。教程用dim=0，只是因为它的张量维度顺序和你默认的不一样。

softmax和argmax的dim选择逻辑是一致的 #

这两个函数的dim参数逻辑完全相同：

• 对softmax来说，dim是你要做归一化的维度，保证该维度上的所有元素和为1；
• 对argmax来说，dim是你要找最大值的维度，返回该维度上最大值的索引。

举个具体例子：
假设logits形状是(3, 5)（3个样本，5个类别）：

• torch.softmax(logits, dim=1)：处理每个样本的5个类别logits，输出形状(3,5)，每行和为1；
• torch.argmax(logits, dim=1)：找每个样本的最大logits类别，输出形状(3)。

如果logits形状是(5, 3)（5个类别，3个样本）：

• torch.softmax(logits, dim=0)：处理每个样本的5个类别logits，输出形状(5,3)，每列和为1；
• torch.argmax(logits, dim=0)：找每个样本的最大logits类别，输出形状(3)。

所以两者的dim选择没有区别，都是跟着类别所在的维度走，只是你可能混淆了不同的张量维度定义场景。

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