tensorflow-gpu运行CNN模型报错ResourceExhaustedError内存分配失败怎么解决

报错原因 #

你遇到的ResourceExhaustedError: failed to allocate memory是典型的GPU显存不足错误。你的输入图像尺寸为512*640，经过三次卷积池化操作后展开为一维向量的长度过大，第一层全连接层初始化权重需要的显存超过了当前GPU的可用剩余显存，所以触发了分配失败报错。之前用CPU训练时调用的是内存，内存容量普遍远大于GPU显存，因此可以正常运行仅速度较慢。

解决方法 #

• 降低输入图像分辨率：将输入的512640尺寸缩小到256320或更小，直接降低特征图展开后的向量长度，大幅减少全连接层参数量
• 替换全连接层前置结构：把Flatten()层替换为tf.keras.layers.GlobalAveragePooling2D()或GlobalMaxPooling2D()，无需将二维特征图完全展开为长向量，可减少90%以上的全连接层参数量，是图像分类任务的常规优化方案
• 减少卷积核数量：适当降低各卷积层的filters数值，比如将最后几层256通道的卷积核调整为128，减少特征图的通道数
• 调低训练批次大小：如果训练时设置的batch_size过高，可以下调到8、4甚至2，降低单步训练需要加载的样本显存占用
• 释放GPU显存：运行程序前关闭其他占用GPU显存的进程，比如其他训练任务、大型3D软件、游戏等，确保空闲显存充足
• 优化损失函数：当前二分类任务使用MSE损失并不适配，替换为sparse_categorical_crossentropy（标签为整数格式）或categorical_crossentropy（标签为onehot编码格式），可加快模型收敛速度，同时降低训练时的计算显存占用

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