介绍

torch.softmax() 函数（实际常用 torch.nn.functional.softmax），这是深度学习多分类任务中最核心的激活函数

softmax函数公式

对于输入：

z = [z_1,z_2,\ldots,z_n]

softmax 计算每个元素的输出：

\text{softmax}(z_i) = \frac{e^{z_i}}{\sum_{j=1}^n e^{z_j}}

\text{j} = 1,2,3, \ldots ,n

特征

• 输出值在 (0, 1) 之间；
• 所有输出值之 和为 1；
• 指数特性会放大数值差异（大值更大，小值更小），适合突出“最可能的类别”。

语法

torch.nn.functional.softmax(input, dim=None, dtype=None)

将任意实数范围的输入张量（通常叫 logits），按指定维度归一化为 0~1 之间的概率分布，且该维度下所有元素的总和为 1。

参数

dim 参数的理解（轴的编号）

张量的维度编号从 0 开始，比如 shape [2,3] 的张量：

• dim=0：垂直方向（行）求和；
• dim=1：水平方向（列）求和。

dim 参数选错的后果

# 错误示例：对 batch 维度（dim=0）归一化（无意义）
wrong_softmax = F.softmax(batch_logits, dim=0)
print("错误dim=0的输出：\n", wrong_softmax)
print("错误求和验证：", wrong_softmax.sum(dim=0))  # tensor([1.0000, 1.0000, 1.0000])
# 后果：每个「类别列」求和为1，而非每个样本的「类别行」求和为1，完全不符合分类逻辑

例子

先通过代码掌握核心用法，注意 dim 参数是最关键的必传参数：

import torch
import torch.nn.functional as F

# 示例1：一维张量（单样本，3个类别）
logits = torch.tensor([1.0, 2.0, 3.0])  # 模型原始输出（logits）
softmax_out = F.softmax(logits, dim=0)

print("softmax输出：", softmax_out)  # tensor([0.0900, 0.2447, 0.6652])
print("求和验证：", softmax_out.sum())  # tensor(1.0)（满足概率分布）

# 示例2：二维张量（批量样本，2个样本×3个类别）
batch_logits = torch.tensor([[1.0, 2.0, 3.0],
                             [4.0, 5.0, 6.0]])
# dim=1：对每个样本的「类别维度」做归一化（核心场景）
batch_softmax = F.softmax(batch_logits, dim=1)

print("批量softmax输出：\n", batch_softmax)
# 输出：
# [[0.0900, 0.2447, 0.6652],  # 第一个样本的类别概率
#  [0.0900, 0.2447, 0.6652]]  # 第二个样本的类别概率
print("每个样本求和：", batch_softmax.sum(dim=1))  # tensor([1.0000, 1.0000])

注意

数值稳定性（内部优化）

直接按公式计算 softmax 容易因 e^x 溢出（比如 logits 为 1000 时，e^1000 超出浮点数范围），PyTorch 的 softmax 内部已做数值稳定化优化：

# 原理：先减去输入维度的最大值，再计算指数
def stable_softmax(x, dim):
    x_max = torch.max(x, dim=dim, keepdim=True)[0]
    x_stable = x - x_max  # 避免指数溢出
    exp_x = torch.exp(x_stable)
    return exp_x / exp_x.sum(dim=dim, keepdim=True)

# 验证溢出场景
overflow_logits = torch.tensor([1000.0, 1001.0, 1002.0])
# PyTorch 官方 softmax 正常输出
print(F.softmax(overflow_logits, dim=0))  # tensor([0.0900, 0.2447, 0.6652])
# 手动实现的稳定版也一致
print(stable_softmax(overflow_logits, dim=0))  # tensor([0.0900, 0.2447, 0.6652])

提示：无需手动做这个优化，PyTorch 已内置实现！

应用场景

多分类任务的完整流程

# 模拟：模型输出 logits → softmax 转概率 → 取最大概率的类别
logits = torch.tensor([[2.5, 1.8, 3.2], [0.5, 4.1, 2.0]])
prob = F.softmax(logits, dim=1)
pred = torch.argmax(prob, dim=1)  # 取概率最大的类别索引

print("类别概率：\n", prob)
print("预测类别：", pred)  # tensor([2, 1])（第一个样本预测类别2，第二个预测类别1）

（2）不要和 CrossEntropyLoss 重复使用

PyTorch 的 nn.CrossEntropyLoss 内部已经包含了 softmax 计算，如果输入是 logits，无需提前做 softmax：

import torch.nn as nn

# 错误：logits → softmax → CrossEntropyLoss（重复计算，结果错误）
logits = torch.tensor([[1.0, 2.0, 3.0]])
label = torch.tensor([2])
prob = F.softmax(logits, dim=1)
loss_wrong = nn.CrossEntropyLoss()(prob, label)

# 正确：直接用 logits 输入 CrossEntropyLoss
loss_right = nn.CrossEntropyLoss()(logits, label)

print("错误损失：", loss_wrong)  # tensor(1.1019)
print("正确损失：", loss_right)  # tensor(0.4076)（符合预期）

（3）softmax 与 log_softmax

如果需要计算 log 概率（比如配合 NLLLoss），优先用 F.log_softmax 而非 torch.log(F.softmax())，数值更稳定：

# 推荐：直接用 log_softmax
log_prob = F.log_softmax(logits, dim=1)
# 不推荐：log(softmax) 可能因数值过小导致精度丢失
log_prob_bad = torch.log(F.softmax(logits, dim=1))

print("log_softmax输出：\n", log_prob)
print("数值一致性：", torch.allclose(log_prob, log_prob_bad))  # True（简单场景下一致）

总结

1. torch.nn.functional.softmax 核心作用是按指定维度将 logits 归一化为和为 1 的概率分布，是多分类任务的标配。
2. dim 参数是关键：批量多分类中必须设为 dim=1（对每个样本的类别维度归一化），选错维度会导致结果无意义。
3. 实战注意：CrossEntropyLoss 内置 softmax，输入无需提前处理；需 log 概率时用 log_softmax 更稳定。

原文出处：/show_1GW2blPk6St3.html