M1 Mac训练EfficientNetB3 CNN大batch时内存不足如何解决

• 该故障属于16GB统一内存M1机型搭配tensorflow-metal运行时的典型内存瓶颈，和数据集加载逻辑、模型结构写法无关。
• 512x512分辨率下训练EfficientNetB3的内存占用远高于常规预估：单张float32格式的输入图像占3MB左右，叠加前向传播中间激活值、反向传播梯度缓存、Metal插件运行时预留内存，batch_size=64时仅训练进程的统一内存占用就会达到12-14GB，加上macOS系统、Jupyter内核、后台进程的基础内存占用，会直接触发内存交换卡死，表现就是停在GPU优化器加载日志无输出、内核终止超时、系统弹出内存不足警告。

所有配置需要写在导入数据集、构建模型的代码之前，否则不会生效。

• 开启GPU内存动态增长，限制TensorFlow的最大内存占用，避免启动时抢占全部可用统一内存：

import tensorflow as tf
gpus = tf.config.list_physical_devices('GPU')
if gpus:
    try:
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
        # 预留6GB内存给系统和其他进程，限制TensorFlow最多使用10GB统一内存
        tf.config.set_logical_device_configuration(
            gpus[0],
            [tf.config.LogicalDeviceConfiguration(memory_limit=10240)]
        )
    except RuntimeError as e:
        print(e)

• 开启混合精度训练，将绝大多数算子的计算精度从float32切换为float16，直接砍半内存占用，同时可调用M1芯片的神经引擎加速计算：

tf.keras.mixed_precision.set_global_policy('mixed_float16')

注意：模型最后输出层如果使用sigmoid或softmax激活，需要将激活逻辑拆分到单独的float32精度层计算，避免数值溢出。

• 分阶段做迁移学习，不要一开始就解冻EfficientNetB3的全部权重：训练初期先冻结EfficientNetB3主干的所有参数，仅训练最后一层全连接分类头，这一步可降低60%以上的梯度缓存内存占用，等训练损失收敛后再逐步解冻顶层网络做微调。
• 调整数据加载管线的参数：不要将prefetch、num_parallel_calls参数直接设为tf.data.AUTOTUNE，将预取buffer大小设为2，并行加载线程数设为4，避免管线提前缓存过多batch占用内存。

• 以上配置全部调整完成后，16GB M1 Air本地运行512分辨率EfficientNetB3训练的稳定batch_size上限可达48-64，完全满足常规训练需求。不需要盲目追求过大batch_size：batch_size超过32后，分类任务的精度提升边际效应极强，配合线性学习率缩放策略，batch=32和batch=128的最终精度差距通常不超过1%。
• 如果确实需要使用128以上的大batch、或者需要解冻全部网络权重做长时间全参数微调，再考虑切换到云算力平台。本地调优后的训练速度和Colab免费版T4显卡差距不大，不需要一开始就迁移训练环境。

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