PyTorch推理后不能直接用raw output当概率,因输出为logits(未归一化分数),须经F.softmax转换为和为1的概率分布;dim=1(或-1)确保按类别维归一化,并需配合model.eval()和torch.no_grad()保障推理正确性。

PyTorch推理后为什么不能直接用 raw output 当概率?
因为模型最后一层(比如 nn.Linear)输出的是 logits —— 未归一化的分数,不是概率。直接取 argmax 可能碰巧对,但 max 值本身没意义,也不能跨样本比较置信度。
必须经过 F.softmax(或 torch.nn.functional.softmax)转换:它把 logits 映射到 [0,1] 区间,且所有类别和为 1。
- 别用
torch.sigmoid—— 那是给多标签(multi-label)用的,输出不加和约束 - 别在训练时已加了
nn.LogSoftmax或nn.CrossEntropyLoss的模型上重复 softmax —— 后者内部已含 softmax + log,再套一层会出错 - 注意维度:默认对最后一维(
dim=-1)做 softmax,分类数必须在该维上
F.softmax 的 dim 参数到底怎么设?
常见错误是传错 dim,导致每行/每列被独立归一化,结果全变成 1 或 nan。
假设你模型输出 shape 是 (batch_size, num_classes)(最常见),那就必须写 F.softmax(logits, dim=1) —— 对每个样本的各类别分数归一化。
让 macOS 自动化像说话一样简单 —— 用代码控制你的 Mac,无需学习 AppleScript 系统要求 macOS 14+ (Sonoma) Python 3.9+ 需要开启:辅助功能 + 屏幕录制权限
立即学习“Python免费学习笔记(深入)”;
-
dim=0:会对 batch 维归一化 → 第 0 类所有样本加起来为 1,完全错误 -
dim=-1等价于dim=1(当只有两维时),可读性稍差但安全 - 如果是 3D 输出(如序列标注:
(batch, seq_len, num_tags)),需明确写dim=-1,避免歧义
import torch import torch.nn.functional as F <p>logits = torch.tensor([[2.1, -1.0, 0.5], [0.3, 1.8, -0.2]]) probs = F.softmax(logits, dim=1) # ✅ 正确:每行独立归一化</p><h1>probs[0] ≈ tensor([0.82, 0.02, 0.16])</h1>
推理时要不要用 with torch.no_grad() 和 .eval()?
要。这两步不是可选优化,而是行为保障。
-
model.eval()关闭 dropout / batch norm 更新,否则推理结果不稳定 -
with torch.no_grad()防止意外构建计算图,节省显存、加速推理,且避免F.softmax被误纳入梯度流 - 漏掉
.eval()在有 dropout 的模型上,可能让同一输入每次输出不同概率分布
model.eval()
with torch.no_grad():
logits = model(x)
probs = F.softmax(logits, dim=1)
softmax 后取 top-k 概率容易忽略的精度问题
直接用 torch.topk(probs, k) 没问题,但若想同时拿到类别索引和概率值,别先 argmax 再查表 —— 有精度损失风险(尤其是 float16 推理时)。
- 正确做法:对
logits用torch.topk得到索引,再用这些索引从probs中 gather —— 避免两次浮点运算累积误差 - 更稳妥:直接对
logits做 topk,然后只对这 k 个 logits 手动 softmax(数值更稳定) - 别用
np.argmax或转 numpy 再处理 —— 破坏 tensor device 一致性,小样本看不出,大 batch 可能报错
复杂点在于:softmax 本身数值稳定(内部做了减 max 技巧),但如果你手动实现或链式操作太多,就容易触发 underflow/overflow。PyTorch 的 F.softmax 已处理好这点,别绕开它。

















