复现论文CNN图像识别模型训练报ResourceExhaustedError(OOM)求优化方案

显存溢出优化方案 #

你的GeForce 940MX仅配备2GB显存，当前模型的特征图尺寸、参数规模超出了显存上限，结合代码给出以下可落地的优化建议，按优先级排序：

1. 优先调整批次大小
   当前批次大小为32，直接下调到8甚至4，无需修改模型结构即可快速验证是否能正常启动训练。
2. 模型结构优化（核心解决论文模型照搬适配问题）

• 初始卷积层无下采样导致特征图尺寸过大：244x244的输入经过两层无步长、无填充的卷积后，特征图尺寸变为239x239，单批次32张对应的特征张量大小接近250MB，反向传播还需存储双倍梯度数据，直接占满显存。可以给前两层卷积添加strides=2参数，或在第一层卷积后新增一个池化层缩小特征图尺寸，也可将第一层卷积的滤波器数量从64下调到16/32。
• 连续卷积无下采样冗余：第二段的三个连续3x3卷积没有下采样操作，特征图尺寸持续冗余，可将其中某一个卷积的步长设为2，或每两个卷积后插入一个小尺寸池化层，进一步压缩特征图。
• 全连接层参数冗余：Flatten后接了5层全连接层，参数占比极高，可砍掉多余的全连接层，最多保留2层全连接即可，同时可降低全连接层的维度，比如将256维下调到128维，128维下调到64维。

3. 训练配置优化

• 开启TensorFlow显存动态分配，在代码最开头加入以下片段，避免TensorFlow初始化时直接占满所有显存：

import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

• 适当缩小输入图像尺寸，可从244x244下调到128x128或160x160，数据集目标特征不大的前提下对精度影响极小，显存占用可下降60%以上。
• 优化器添加梯度裁剪参数，减少反向传播时的梯度显存占用。

4. 修复代码逻辑错误（消除额外冗余显存占用）

• 输出层配置与损失不匹配：当前输出层为Dense(1, activation='softmax')配合分类交叉熵损失存在逻辑错误，二分类任务请将输出层激活改为sigmoid，损失换为BinaryCrossentropy；多分类任务请将输出层神经元数量改为对应类别数。
• 损失参数from_logits=True仅适用于输出层无激活函数的场景，你当前输出层已加softmax激活，该参数需改为False，否则会产生多余计算消耗显存，同时导致训练精度异常。

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