PyTorch复现Keras路透社多分类模型效果差异排查

我正在学习《Deep Learning with Python (Keras)》一书，为加深理解将书中示例改写为PyTorch代码，多数示例效果与Keras版本相当，但在第106页的路透社多分类任务中遇到问题。

书中Keras代码 #

keras_model = keras.Sequential([
    layers.Dense(64,activation='relu'),
    layers.Dense(64,activation='relu'),
    layers.Dense(46,activation='softmax'),
])

keras_model.compile(
    optimizer = 'rmsprop',
    loss = 'sparse_categorical_crossentropy',
    metrics = ['accuracy']
)

hist = keras_model.fit(
    partial_train_xs,
    partial_train_ys,
    epochs=20,
    batch_size=512,
    validation_data=[val_xs,val_ys]
)

我尝试复现的PyTorch代码 #

model = nn.Sequential(
          nn.Linear(10_000,64),
          nn.ReLU(),
          nn.Linear(64,64),
          nn.ReLU(),
          nn.Linear(64,46),
          nn.Softmax()
        )

def compute_val_loss(model,xs,ys):
    preds = model(xs)
    return(F.cross_entropy(preds,ys)).item()

def compute_accuracy(model,xs,ys):
    preds = model(xs)
    acc = (preds.argmax(dim=1) == ys).sum() / len(preds)
    return acc.item()

def train_loop(model,xs,ys,epochs=20,lr=1e-3,opt=torch.optim.RMSprop,
               batch_size=512,loss_func=F.cross_entropy):
    opt = opt(model.parameters(),lr=lr)
    losses = []
    for i in range(epochs):
        epoch_loss = []
        for b in range(0,len(xs),batch_size):
            xbatch = xs[b:b+batch_size]
            ybatch = ys[b:b+batch_size]            
            logits = model(xbatch)
            loss = loss_func(logits,ybatch)
            model.zero_grad()
            loss.backward()
            opt.step()
            epoch_loss.append(loss.item())
        losses.append([i,sum(epoch_loss)/len(epoch_loss)])
        print(loss.item())
    return losses

数据采用multihot编码处理词序列，例如词汇表为10k时，输入是10k维向量，出现单词对应索引位为1。

问题描述 #

PyTorch版本与Keras版本结果差异显著：

• Keras版本20轮后训练损失极低，验证准确率约80%；
• PyTorch版本训练损失从3.4缓慢降至3.1，Keras训练1轮后的损失（2.6）就已低于该值；
• PyTorch准确率虽有提升但始终落后，且呈阶梯状（如图），即使训练150轮，验证准确率峰值仅约75%。

调整RMSProp参数无明显改善，请问是代码有误还是合理差异？

问题根源与修正方案 #

你的PyTorch代码存在两处关键错误，直接导致训练效率和效果远低于Keras版本：

1. Softmax层与CrossEntropyLoss的冲突 #

PyTorch的F.cross_entropy内置了LogSoftmax和NLLLoss的计算逻辑，但你在模型最后额外添加了nn.Softmax()层，会引发两个问题：

• 数值稳定性问题：Softmax输出的概率值接近0或1时，取对数会出现梯度消失；
• 损失计算逻辑错误：相当于对已经经过Softmax的结果再做一次LogSoftmax，完全偏离正确的损失计算路径。

修正：移除模型中的nn.Softmax()层，让最后一层Linear直接输出logits即可。

2. 优化器参数与训练流程未对齐Keras #

• Keras的RMSprop默认动量为0，而PyTorch的torch.optim.RMSprop默认动量为0.9，这会明显影响收敛速度，需手动设置momentum=0对齐；
• 原训练循环未在每轮结束后计算验证集指标，无法监控泛化能力，也无法和Keras的训练日志对比；
• 评估时未切换模型到eval()模式，也未关闭梯度计算，会浪费计算资源并可能影响结果。

修正后的PyTorch代码示例 #

import torch
import torch.nn as nn
import torch.nn.functional as F

# 移除Softmax层，直接输出logits
model = nn.Sequential(
    nn.Linear(10_000, 64),
    nn.ReLU(),
    nn.Linear(64, 64),
    nn.ReLU(),
    nn.Linear(64, 46)
)

def compute_val_loss(model, xs, ys):
    model.eval()  # 切换到评估模式
    with torch.no_grad():  # 关闭梯度计算，节省资源
        preds = model(xs)
        loss = F.cross_entropy(preds, ys).item()
    model.train()  # 切回训练模式
    return loss

def compute_accuracy(model, xs, ys):
    model.eval()
    with torch.no_grad():
        preds = model(xs)
        acc = (preds.argmax(dim=1) == ys).sum() / len(preds)
    model.train()
    return acc.item()

def train_loop(model, train_xs, train_ys, val_xs, val_ys, epochs=20, lr=1e-3, opt=torch.optim.RMSprop,
               batch_size=512, loss_func=F.cross_entropy):
    # 对齐Keras的RMSprop参数：momentum=0
    opt = opt(model.parameters(), lr=lr, momentum=0)
    train_losses = []
    val_losses = []
    val_accs = []
    
    for i in range(epochs):
        epoch_loss = []
        model.train()
        for b in range(0, len(train_xs), batch_size):
            xbatch = train_xs[b:b+batch_size]
            ybatch = train_ys[b:b+batch_size]            
            logits = model(xbatch)
            loss = loss_func(logits, ybatch)
            
            opt.zero_grad()  # 用优化器清梯度更规范
            loss.backward()
            opt.step()
            
            epoch_loss.append(loss.item())
        
        # 计算并记录本轮训练/验证指标
        avg_train_loss = sum(epoch_loss)/len(epoch_loss)
        train_losses.append(avg_train_loss)
        
        avg_val_loss = compute_val_loss(model, val_xs, val_ys)
        val_acc = compute_accuracy(model, val_xs, val_ys)
        val_losses.append(avg_val_loss)
        val_accs.append(val_acc)
        
        print(f"Epoch {i+1}: Train Loss = {avg_train_loss:.4f}, Val Loss = {avg_val_loss:.4f}, Val Acc = {val_acc:.4f}")
    
    return train_losses, val_losses, val_accs

额外建议 #

• 确保输入数据类型对齐：Keras默认用float32，PyTorch需保证张量为torch.float32，若用GPU则模型和数据需同步移至GPU；
• 权重初始化：Keras和PyTorch的全连接层默认初始化逻辑一致（Xavier/Glorot初始化），若需完全匹配可手动指定初始化方式。

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