TensorFlow模型定义阶段触发GPU ResourceExhaustedError报错

这个报错和训练阶段的batch size没有任何关系，触发点是模型初始化权重时单次申请的连续显存量过大，剩余显存碎片化无法满足分配要求，和报错栈里调用random_uniform生成初始化权重的逻辑完全对应。
算下你当前网络的参数量就清楚了：
输入尺寸4003003，两层same卷积后接步长2、核大小3的valid池化，输出特征图尺寸约为19914964，拉平后维度接近190万。后面接512单元的Dense层时，权重矩阵形状为(1897664, 512)，单这一个float32格式的权重张量就要占约3.7GB显存——和你日志里记录的MaxAllocSize: 3886415872完全匹配。
你当前GPU总显存约10GB，已经占用7.7GB，剩余的2.3GB显存全是碎块（日志里LargestFreeBlock: 0就是直接证据，没有连续的大块空闲显存），根本塞不下3.7GB的连续张量申请，直接触发OOM。
至于nvidia-smi显示显存没占满很正常：nvidia-smi只统计进程总显存占用量，不会展示显存内部分布，连续大块显存不足时，就算总剩余显存数值够，也会分配失败。

• 核心修复：调整网络结构，从根源降低全连接层的参数量
  • 不要在仅做过一次步长2池化的高分辨率特征图后直接Flatten接全连接层。你当前输入分辨率400300，仅一次池化后特征图尺寸还是太大，全连接参数量会爆炸。可以新增2-3组「卷积+步长2池化」的组合，把特征图分辨率压到2518量级再拉平，全连接层参数量会直接降到原来的1%量级。
  • 用全局平均池化替代Flatten+大全连接的组合：直接替换Flatten层为layers.GlobalAveragePooling2D()，该层会对每个通道的整张特征图求均值，输出维度等于当前通道数，不需要存海量全连接权重，是图像分类任务的常规操作，能直接把全连接层参数量降两个数量级。
  • 如果必须保留全连接结构，可以在池化后加1*1卷积先压缩通道数，再做Flatten，同样能大幅降低全连接层的输入维度。
• 显存配置优化，减少碎片
  • 把GPU显存动态增长的配置代码放到所有TensorFlow、Keras相关导入语句的最前面，如果先导入Keras再开显存增长，配置不会生效，框架启动时会直接预留大量显存，产生更多碎片。
  • 可以直接设置环境变量TF_FORCE_GPU_ALLOW_GROWTH=true，生效优先级高于代码层面的set_memory_growth调用，能进一步减少框架启动时的无效显存预留。
  • 初始化模型前先调用tf.keras.backend.clear_session()，清空之前运行残留的计算图和显存占用，释放被占住的大块显存。
• 快速验证方法：如果想快速确认根因，可以临时把输入分辨率降到100*75，或者把512单元的Dense层改成32单元，如果模型能正常初始化，就可以100%确定是全连接层参数量过大导致的问题。

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