M1 Mac训练TensorFlow模型时损失异常暴涨问题求助

问题背景 #

在M1 MacBook Pro上训练图像分类模型时，损失值出现异常暴涨（最高达200万亿），训练精度持续徘徊在随机猜测水平。但相同模型、相同Python版本在WSL2 Ubuntu环境中可正常训练，损失稳步降至0.05左右。更换其他模型（如CIFAR-10简单模型）仍存在类似问题，已尝试重启设备、重装TensorFlow，问题未解决。

环境信息 #

• Python 3.11.4
• TensorFlow 2.13.0
• tensorflow-metal 1.0.1

模型结构（示例） #

模型包含多组Conv2D+MaxPool2D、Dropout层，以及全连接层，最终用softmax输出二分类结果，损失函数为categorical_crossentropy，优化器为Adam（学习率0.001）。

可能原因及解决方案 #

1. tensorflow-metal后端数值稳定性问题 #

M1/M2芯片的GPU浮点运算特性与x86 GPU存在差异，tensorflow-metal 1.0.1等早期版本对MaxPool2D、Conv2D的梯度计算支持不完善，易引发梯度爆炸，导致损失暴涨。

解决措施：

• 降低学习率：将Adam的学习率从0.001下调至0.0001或0.00001，减少单次梯度更新的幅度，避免数值溢出。
• 临时切换CPU验证：添加代码禁用GPU，确认是否为GPU后端问题：

import tensorflow as tf
tf.config.set_visible_devices([], 'GPU')

若CPU训练损失正常，建议升级tensorflow-metal至最新稳定版，同时匹配对应版本的TensorFlow（遵循官方兼容表）。

2. 优化器金属后端适配bug #

尽管日志显示已自动回退到legacy.Adam，但新版TensorFlow优化器在金属后端的梯度计算仍可能存在异常，导致梯度累积失控。

解决措施：

• 显式指定legacy优化器：直接调用tf.keras.optimizers.legacy.Adam，避免自动回退的潜在问题：

from tensorflow.keras.optimizers.legacy import Adam
model.compile(loss="categorical_crossentropy", optimizer=Adam(learning_rate=0.0001), metrics=["accuracy"])

• 尝试其他优化器：改用SGD（带动量）替代Adam，验证是否能稳定训练：

model.compile(loss="categorical_crossentropy", optimizer=tf.keras.optimizers.SGD(learning_rate=0.001, momentum=0.9), metrics=["accuracy"])

3. 输入数据预处理不一致 #

Mac与WSL2环境中图像生成器的预处理流程可能存在差异（如像素值归一化范围），若输入数据未正确缩放，会导致模型输入值过大，引发损失爆炸。

解决措施：

• 统一数据预处理：确保两个环境中都将图像像素值缩放到[0,1]范围：

from tensorflow.keras.preprocessing.image import ImageDataGenerator

train_datagen = ImageDataGenerator(rescale=1./255)
train_generator = train_datagen.flow_from_directory(...)
valid_generator = train_datagen.flow_from_directory(...)

• 验证数据分布：检查两个环境中训练数据的均值、方差，确保数据分布完全一致。

4. Python与TensorFlow版本兼容性问题 #

Python 3.11.x与TensorFlow 2.13.0的组合在M1芯片上可能存在未适配的底层问题，影响数值运算稳定性。

解决措施：

• 降级Python版本：切换到TensorFlow官方推荐的Python 3.10.x版本，重新搭建环境后测试。
• 匹配兼容版本：根据官方文档，选择TensorFlow与tensorflow-metal的兼容组合（如TensorFlow 2.16.x + tensorflow-metal 1.2.0）。

验证流程 #

1. 先切换到CPU训练，确认损失是否恢复正常，排查GPU后端问题。
2. 调整学习率并显式使用legacy Adam，重新训练观察损失变化。
3. 核对数据预处理代码，确保与WSL2环境完全一致。
4. 若以上无效，尝试降级Python或升级tensorflow-metal版本。

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