PyTorch替换log_softmax为F.cross_entropy后结果异常问题定位

输出异常的核心原因是模型输出格式和损失函数要求的输入格式不匹配，代码执行顺序和教程的重构节奏完全错位。

教程原本的逻辑匹配关系 #

教程从自定义实现过渡到内置组件的过程，严格遵循“组件同步替换”的原则：

• 初始自定义阶段的整套逻辑是自洽的：
  • 模型通过log_softmax把线性计算结果转换为归一化后的对数概率值
  • 配套的自定义nll（负对数似然损失），专门接收对数概率作为输入，计算对应类别的损失均值
• 重构后阶段的逻辑是PyTorch对两个步骤的封装：
  • 模型简化为直接输出未经处理的线性计算值（也就是常说的logits），不再手动做log_softmax
  • 损失函数替换为内置的F.cross_entropy——这个函数的内部已经合并了log_softmax计算和负对数似然计算的步骤，专门接收logits作为输入
  • 这套组合的计算结果和最初的“手动log_softmax + 自定义nll”组合完全等价

代码执行的实际错位点 #

按代码从上到下的执行顺序，操作出现了两个关键错位：

1. 虽然把模型改成了直接输出logits，但紧接着又重新定义了nll函数，并且把loss_func重新赋值回了这个要求“对数概率输入”的自定义损失：

   def nll(input, target):
       return -input[range(target.shape[0]), target].mean()
   loss_func = nll
   

2. 后续的整个训练循环，全程都在用这个和新模型不匹配的损失计算梯度：
   把未归一化的原始logits直接传给只接受对数概率的nll函数，相当于跳过了log_softmax的归一化步骤，直接对原始分数取负计算平均，自然会得到偏离合理范围的loss值（负数、超大正数都是这类错误计算的典型结果），同时梯度方向完全错误，训练后的权重也处于异常状态。
3. 等两个epoch的训练全部跑完，才在最后打印结果前，把loss_func换成了正确的F.cross_entropy，但此时模型权重已经被前面错误的训练流程破坏了。

正确的修正步骤 #

• 删掉中间重复定义nll、重新赋值loss_func = nll的代码块
• 在把模型改成直接输出logits之后，立刻将loss_func赋值为F.cross_entropy，保证训练全程模型和损失函数的输入输出要求匹配
• 训练完成后再做结果打印，就能复现教程中两个实现结果完全一致的效果

修正后的核心代码片段参考：

import torch.nn.functional as F

# 模型直接输出logits，不再手动调用log_softmax
def model(xb):
    return xb @ weights + bias

# 损失函数立即匹配为内置交叉熵，不要后续重新赋值回自定义nll
loss_func = F.cross_entropy

# 后续训练循环全程使用匹配好的损失函数
lr = 0.5
epochs = 2
for epoch in range(epochs):
    for i in range((n - 1) // bs + 1):
        start_i = i * bs
        end_i = start_i + bs
        xb = x_train[start_i:end_i]
        yb = y_train[start_i:end_i]
        pred = model(xb)
        loss = loss_func(pred, yb)

        loss.backward()
        with torch.no_grad():
            weights -= weights.grad * lr
            bias -= bias.grad * lr
            weights.grad.zero_()
            bias.grad.zero_()

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