M1 Mac上TensorFlow函数式CNN分类器AUC徘徊0.5，Colab训练正常

问题背景 #

使用TensorFlow构建二分类图像CNN模型，损失函数为binary cross-entropy，评估指标为ROC AUC。同一Jupyter Notebook在两个平台运行结果完全不同：

• Google Colab：训练集指标2-3个epoch后达到峰值，随后严重过拟合
• M1 MacBook Pro（GPU）：训练集ROC AUC始终在0.5左右波动，验证集AUC固定为0.5，模型完全无法学习

已设置随机种子保证可复现性，且M1上其他模型架构训练同一数据集可正常运行。M1环境通过pip安装tensorflow、tensorflow-macos、tensorflow-metal配置GPU支持，曾出现模型编译卡顿问题。

代码重现 #

import numpy as np
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras.preprocessing import image_dataset_from_directory
import tensorflow as tf
import matplotlib.pyplot as plt
import datetime

POSITIVE_COUNT = 393
NEGATIVE_COUNT = 400666
TOTAL_COUNT = 401059
IMG_SIZE = (224, 224)
BATCH_SIZE = 32
SEED = 12
pos_dataset = image_dataset_from_directory('nvm/1', labels=[1] * POSITIVE_COUNT,
                                      seed=SEED, image_size=IMG_SIZE, batch_size=BATCH_SIZE, validation_split=0.3, subset='both')
neg_dataset = image_dataset_from_directory('nvm/0', labels=[0] * NEGATIVE_COUNT,
                                      seed=SEED, image_size=IMG_SIZE, batch_size=BATCH_SIZE, validation_split=0.3, subset='both')

train_ds = tf.data.Dataset.sample_from_datasets(
    [neg_dataset[0], pos_dataset[0].repeat(1019)], seed=SEED, stop_on_empty_dataset=False, rerandomize_each_iteration=True
)
val_ds = tf.data.Dataset.sample_from_datasets(
    [neg_dataset[1], pos_dataset[1]], seed=SEED, stop_on_empty_dataset=False
)

augmentation = keras.Sequential([
    layers.RandomRotation((-0.3, 0.3), seed=SEED),
    layers.RandomTranslation((-0.1, 0.1), (-0.1, 0.1), seed=SEED),
    layers.RandomZoom(0.2, seed=SEED),
    layers.RandomFlip(mode='horizontal', seed=SEED)
])

def create_model():
    inputs = keras.Input(shape=(*IMG_SIZE, 3))
    x = augmentation(inputs)
    x = layers.Rescaling(1./255)(inputs)
    x = layers.Conv2D(filters=32, kernel_size=(3, 3), activation="relu")(x)
    x = layers.MaxPooling2D(pool_size=(2, 2))(x)
    x = layers.Conv2D(filters=64, kernel_size=(3, 3), activation="relu")(x)
    x = layers.MaxPooling2D(pool_size=(2, 2))(x)
    x = layers.Conv2D(filters=64, kernel_size=(3, 3), activation="relu")(x)
    x = layers.MaxPooling2D(pool_size=(2, 2))(x)
    x = layers.Conv2D(filters=128, kernel_size=(3, 3), activation="relu")(x)
    x = layers.MaxPooling2D(pool_size=(2, 2))(x)
    x = layers.Dropout(0.5)(x)
    x = layers.Flatten()(x)
    x = layers.Dense(64, activation="relu")(x)
    x = layers.Dropout(0.3)(x)
    outputs = layers.Dense(1, activation="sigmoid")(x)
    model = keras.Model(inputs=inputs, outputs=outputs)
    optimizer = tf.keras.optimizers.Adam(learning_rate=0.001, ema_momentum=0.99)
    model.compile(loss="binary_crossentropy", optimizer=optimizer, metrics=[keras.metrics.AUC()])
    return model

log_dir = "logs/fit/" + 'exp2_' + datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=log_dir, histogram_freq=1)

model_concat = create_model()
history = model_concat.fit(train_ds, epochs=20, validation_data=val_ds, callbacks=[tensorboard_callback])

原因分析及解决方案 #

1. 修复代码中的核心bug #

代码中数据增强层的结果被完全丢弃，这是导致模型行为异常的关键问题：

# 原错误代码
x = augmentation(inputs)  # 执行增强，但结果未使用
x = layers.Rescaling(1./255)(inputs)  # 直接使用原始输入，增强失效

# 修正后
x = augmentation(inputs)
x = layers.Rescaling(1./255)(x)  # 使用增强后的张量继续处理

该bug在两个平台都存在，但Colab上数据集本身可能因为类别不平衡（正样本极少），即使无增强也能快速过拟合，而M1平台叠加其他兼容性问题导致完全无法学习。

2. M1 TensorFlow-metal兼容性适配 #

禁用EMA优化器参数

TensorFlow-metal对Adam优化器的ema_momentum参数支持存在兼容性问题，尝试移除该参数使用标准Adam：

optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)

替换类别不平衡处理方式

使用class_weight替代数据集重复的方式，避免repeat与sample_from_datasets在M1上的交互异常：

# 计算类别权重
weight_for_0 = (1 / NEGATIVE_COUNT) * (TOTAL_COUNT) / 2.0
weight_for_1 = (1 / POSITIVE_COUNT) * (TOTAL_COUNT) / 2.0
class_weight = {0: weight_for_0, 1: weight_for_1}

# 修改训练集构建，无需重复正样本
train_ds = tf.data.Dataset.sample_from_datasets(
    [neg_dataset[0], pos_dataset[0]], seed=SEED, weights=[0.5, 0.5], 
    stop_on_empty_dataset=False, rerandomize_each_iteration=True
)

# 训练时传入类别权重
history = model_concat.fit(
    train_ds, epochs=20, validation_data=val_ds, 
    callbacks=[tensorboard_callback], class_weight=class_weight
)

验证GPU兼容性问题

临时禁用GPU，改用CPU训练，确认是否为TensorFlow-metal的问题：

tf.config.set_visible_devices([], 'GPU')

若CPU训练正常，更新到最新版本的TensorFlow组件：

pip install --upgrade tensorflow-macos tensorflow-metal

检查模型初始化

M1上部分层的初始化可能存在异常，手动指定初始化器并验证权重分布：

# 修改Conv2D层使用he_normal初始化
layers.Conv2D(filters=32, kernel_size=(3,3), activation="relu", kernel_initializer="he_normal")

# 打印初始权重均值，确认分布正常
model = create_model()
for layer in model.layers:
    if hasattr(layer, 'kernel'):
        print(f"{layer.name}: {layer.kernel.numpy().mean():.6f}")

3. 强化随机种子一致性 #

确保所有随机环节的种子覆盖，消除平台间随机数实现差异：

tf.random.set_seed(SEED)
np.random.seed(SEED)
import random
random.seed(SEED)

# 可选：启用tf.data调试模式，排查数据管道问题
tf.data.experimental.enable_debug_mode()

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