TensorFlow与PyTorch LSTM行为差异原因排查求助

1. AdamW优化器默认参数差异 #

这是最可能的直接原因：

• PyTorch的torch.optim.AdamW默认weight_decay=0.01，而TensorFlow的tf.keras.optimizers.AdamW默认weight_decay=0.0。
• 你提到两者都用AdamW但只指定了学习率，未设置weight decay。PyTorch端的模型会被施加0.01的权重衰减，对于金融时序数据，这会导致权重更新被过度正则化，无法有效拟合数据，最终表现为损失居高不下。

2. LSTM偏置初始化差异 #

PyTorch与TensorFlow的LSTM遗忘门偏置初始化策略不同：

• PyTorch的nn.LSTM默认将遗忘门的偏置初始化为1（帮助模型初始时保留更多历史记忆），而TensorFlow Keras的LSTM默认遗忘门偏置为0。
• 这种初始化差异会导致模型初始阶段的梯度流动特性不同，对于某些时序数据，PyTorch的初始设置可能延缓模型收敛，甚至在结合weight decay时完全阻碍学习。

3. 数据精度与梯度数值稳定性 #

你提到用float32时PyTorch后端的Keras也无法学习，说明float32精度下的梯度数值稳定性是另一个影响因素：

• TensorFlow在float32模式下默认会做一些数值优化（如自动混合精度的隐式处理），而PyTorch需要显式开启相关功能才能获得类似的稳定性。
• 金融数据的数值波动可能导致float32下梯度消失/爆炸，而PyTorch默认的float32训练没有这些优化，进一步加剧学习困难。

验证与解决步骤 #

1. 统一AdamW参数：
   在PyTorch中初始化AdamW时显式设置weight_decay=0，和TensorFlow保持一致：

   optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.0)
   

   这是最优先要验证的点，大概率能解决损失不下降的问题。

2. 调整LSTM偏置初始化：
   如果统一weight decay后仍有问题，可以手动重置PyTorch LSTM的遗忘门偏置为0，对齐TensorFlow的初始化：

   # 获取LSTM的偏置参数，遗忘门偏置在每个layer的bias中占1/4
   for name, param in model.lstm.named_parameters():
       if 'bias' in name:
           # 遗忘门偏置是bias的第hidden_size到2*hidden_size部分
           n = param.size(0)
           start, end = n//4, n//2
           param.data[start:end].fill_(0.)
   

3. 开启自动混合精度训练：
   针对float32精度的稳定性问题，在PyTorch中启用自动混合精度：

   scaler = torch.cuda.amp.GradScaler()
   for epoch in range(epochs):
       model.train()
       for x, y in dataloader:
           optimizer.zero_grad()
           with torch.cuda.amp.autocast():
               preds = model(x)
               loss = nn.L1Loss()(preds, y)
           scaler.scale(loss).backward()
           scaler.step(optimizer)
           scaler.update()
   

4. 确认输入数据一致性：
   确保PyTorch和TensorFlow使用的输入数据完全一致（包括归一化方式、数据类型、维度顺序），比如检查是否存在数据尺度差异、预处理步骤不一致等问题。

内容的提问来源于stack exchange，提问作者Benjamin van Heerden