Keras函数式API定义模型触发ResourceExhaustedError显存不足报错

核心问题排查 #

1. 代码逻辑错误

• 存在变量覆盖问题：定义covn03层后，你重复对acti02、pool02变量赋值，导致covn03层完全没有参与计算，同时计算图中残留无用张量，额外占用显存。
• 输入定义错误：model_one和model_two共享同一个Input层实例，最终合并模型时你把inputs设置为两个输入的列表，会导致框架重复加载两份输入张量，显存占用直接翻倍。

2. 显存溢出根本原因

你将两个子模型的特征图拼接后直接做Flatten，参数量会出现爆炸式增长。以输入尺寸2242243为例，经过3次2倍下采样后特征图尺寸为28*28，两个子模型各输出64通道，拼接后总通道128，Flatten后的向量长度为28*28*128=100352，仅第一层Dense(256)的参数就超过2500万，再加上反向传播需要存储的梯度值，batch size稍大就会占满RTX2060的6G显存。
你之前在Mac上用CPU跑可以正常运行，是因为CPU版本调用的是系统内存，8G系统内存比2060的显存容量更大，所以不会触发内存不足报错。

解决方法 #

1. 修复代码错误
   • 修正变量命名，删掉无用的covn03层，或者将对应的激活、池化层变量改为acti03、pool03，保证子模型计算逻辑正确。
   • 如果是单输入场景（同一份输入喂给两个子模型），最终模型的inputs只需要填写1个输入实例即可，不要写成列表；如果是双输入场景，需要给两个子模型分别定义独立的Input层，不要共享同一个实例。
2. 大幅降低参数量
   把concat层后的Flatten替换为全局平均池化，直接将特征图压缩为和通道数相等的向量，参数量可以降到原来的1%：

   concat_feature_layer = layers.concatenate([model_one.output, model_two.output])
   gap_layer = layers.GlobalAveragePooling2D()(concat_feature_layer)
   fully_connected_dense_big = layers.Dense(256, activation='relu')(gap_layer)
   

3. 优化显存分配
   在代码最开头加入以下配置，开启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)
   

4. 调小训练batch size
   如果修改后仍报错，把训练时的batch size降到8、4甚至2即可。

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