GPU训练半轮后出现Loss NaN但CPU无异常的问题求助

核心现象 #

• GTX 1080 GPU训练时，loss与val-loss波动剧烈，快速出现Loss NaN；大幅降低学习率可避免NaN，但模型性能极差
• 同初始化参数（加载同一模型文件）下，CPU训练表现完全正常：高学习率时loss稳定下降；即使与GPU用相同学习率，相同轮数下CPU的loss/val-loss更低
• 云端A10 GPU运行完全相同的模型与数据，无任何异常

模型与训练配置 #

• 模型结构：3个并行LSTM分支（输入为32步×12特征），每个分支包含512单元（return_sequences=True）+128单元LSTM层；分支输出拼接后接2层256单元Dense层（SELU激活），最后输出1维回归结果
• 正则化策略：LSTM层添加小量L2正则化（0.00001），Dense层添加0.05 Dropout；有无正则化/Dropout的场景均已测试
• 损失函数：代码配置为MAE，同时尝试过MSE
• 已尝试的调试动作：各类正则化手段、梯度裁剪（clip value/norm）、降低学习率、GPU烤机测试（无硬件报错）、更换模型架构（含无LSTM的纯Dense模型）、调整批次大小（极小批次无效）

环境信息 #

• GPU：GTX 1080
• 软件栈：TensorFlow 2.13、cuDNN 8.6、Nvidia驱动535.86.05
• 已重装Ubuntu并严格遵循TensorFlow官方安装指南，问题未解决

完整模型代码 #

num_properties = 12
dropout_rate = 0.05
l2_reg_value = 0.00001

input_branch1 = Input(shape=(32, num_properties))
input_branch2 = Input(shape=(32, num_properties))
input_branch3 = Input(shape=(32, num_properties))

lstm_branch1 = LSTM(512, return_sequences=True, kernel_regularizer=l2(l2_reg_value))(input_branch1)
lstm_branch1 = LSTM(128, kernel_regularizer=l2(l2_reg_value))(lstm_branch1)

lstm_branch2 = LSTM(512, return_sequences=True, kernel_regularizer=l2(l2_reg_value))(input_branch2)
lstm_branch2 = LSTM(128, kernel_regularizer=l2(l2_reg_value))(lstm_branch2)

lstm_branch3 = LSTM(512, return_sequences=True, kernel_regularizer=l2(l2_reg_value))(input_branch3)
lstm_branch3 = LSTM(128, kernel_regularizer=l2(l2_reg_value))(lstm_branch3)

concatenated = concatenate([lstm_branch1, lstm_branch2, lstm_branch3])

dense = Dense(256, activation='selu')(concatenated)
dense = Dropout(dropout_rate)(dense)
dense = Dense(256, activation='selu')(dense)
dense = Dropout(dropout_rate)(dense)

output = Dense(1, activation='linear')(dense)

model = Model(inputs=[input_branch1, input_branch2, input_branch3], outputs=output)

model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.0001, clipnorm=1), loss="mae")

具体解决建议 #

1. 降级TensorFlow版本：GTX 1080属于Pascal架构，TensorFlow 2.13对老架构的兼容性可能存在瑕疵，建议降级到2.10或2.11版本（这两个版本对Pascal GPU的支持更成熟稳定）
2. 匹配cuDNN与TensorFlow版本：TensorFlow 2.13官方推荐的cuDNN版本为8.9.0，当前使用的cuDNN 8.6版本不匹配，可能导致LSTM计算时出现浮点异常，建议更换为对应版本的cuDNN
3. 禁用cuDNN加速LSTM：尝试强制关闭cuDNN对LSTM的优化，验证是否为cuDNN适配问题，可在代码开头添加：

   # 禁用cuDNN LSTM加速（部分TF版本支持）
   from tensorflow.keras.layers import LSTM
   tf.keras.layers.LSTM = lambda *args, **kwargs: LSTM(*args, **kwargs, use_cudnn=False)
   

   或者开启内存增长模式避免显存分配问题：

   gpus = tf.config.list_physical_devices('GPU')
   if gpus:
       try:
           for gpu in gpus:
               tf.config.experimental.set_memory_growth(gpu, True)
       except RuntimeError as e:
           print(e)
   

4. 严格检查数据标准化：确认输入数据是否存在极端值，GPU的浮点计算精度（FP32）与CPU可能存在差异，极端值容易引发NaN，建议将输入数据缩放到[-1,1]范围或做Z-score标准化
5. 调整优化器参数：尝试增大Adam优化器的epsilon参数（默认1e-7）到1e-5，避免因浮点精度问题导致梯度计算异常；也可更换为SGD+动量优化器测试
6. GPU硬件隐性故障排查：运行简单的LSTM测试代码（比如小批量数据训练单LSTM层），验证GPU上的LSTM计算是否正常，排除硬件隐性问题

内容的提问来源于stack exchange，提问作者Christian Reinholdt Andersen