CNN训练触发OOM(ResourceExhaustedError)如何缩减模型参数量

你当前显存溢出的核心诱因是架构设计缺陷：仅对224x224输入做了1次2倍下采样的最大池化，得到112x112分辨率、128通道的特征图后直接展平，得到长度为1605632的一维向量，后续接128单元全连接层时单层参数量就超过2亿，远超过常规消费级GPU的显存承载能力。
注：以下所有调整方案均保留224x224的输入分辨率，完全适配后续使用224分辨率输入预训练ViT的需求。

方案1：堆叠卷积+池化单元逐步压缩特征图尺寸 #

这是经典CNN的标准设计逻辑：每堆叠1-2层卷积后接一次2倍下采样的池化层，逐步把特征图的空间分辨率压低，再展平接全连接层，从根源上降低展平后的向量长度。
参考修改后的模型结构代码：

input_shape = (224, 224, 3)
cnn_model = tf.keras.Sequential([
    # 第一个卷积块 224→112
    Conv2D(32, (5,5), padding = 'same', activation = 'relu', input_shape = input_shape),
    MaxPooling2D(2,2),
    # 第二个卷积块 112→56
    Conv2D(64, (3,3), padding = 'same', activation = 'relu'),
    MaxPooling2D(2,2),
    # 第三个卷积块 56→28
    Conv2D(128, (3,3), padding = 'same', activation = 'relu'),
    MaxPooling2D(2,2),
    # 第四个卷积块 28→14
    Conv2D(128, (3,3), padding = 'same', activation = 'relu'),
    MaxPooling2D(2,2),
    # 此时特征图尺寸为14*14*128=25088，展平后维度比原模型小64倍
    Flatten(),
    Dense(128, activation = 'relu'),
    Dense(1, activation = "sigmoid"),                       
])

调整后全连接层参数量从2亿降到300万级别，显存占用会大幅下降。

方案2：用全局平均池化替代Flatten层 #

如果想要进一步压缩参数量，可以直接去掉Flatten层，替换为GlobalAveragePooling2D层。该操作会对每个通道的整个空间特征求平均值，无论输入特征图的空间分辨率多大，输出维度永远等于卷积层的通道数，完全不会出现百万级长度的展平向量。
参考修改片段：

# 替换原模型中的Flatten()即可
MaxPooling2D(2,2),
GlobalAveragePooling2D(), # 输出维度直接为128
Dense(128, activation = 'relu'),
Dense(1, activation = "sigmoid"),

这种改法下，全连接层参数量仅为128*128 + 128 = 16512，参数量比原模型低6个数量级，哪怕是入门级GPU也能轻松跑通，同时全局平均池化本身还自带一定的防过拟合效果。

辅助优化手段（非必须） #

• 训练时如果仍有显存压力，可以适当调小batch size，不需要修改模型结构
• 可以在卷积层、全连接层之间加入BatchNormalization或低比例Dropout，在几乎不增加显存占用的前提下降低过拟合风险

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