TensorFlow 2.5训练Faster R-CNN触发GPU OOM显存耗尽报错咨询

• 核查TensorFlow显存分配策略：默认TensorFlow会抢占几乎所有可用显存，即便实际计算不需要对应容量。报错信息中显示可用显存上限仅为约5.8GB（6269894656字节），未达到RTX2070的8GB标称显存，说明存在其他进程占用显存，可执行nvidia-smi命令核查后台占用显存的进程，关闭无关的python训练进程、浏览器硬件加速、重型桌面渲染任务等释放显存。
• 核查模型与配置匹配度：Faster R-CNN ResNet101 V1本身参数量较大，640x640输入分辨率下单张推理就需要接近2GB显存，训练阶段需要额外存储梯度、优化器状态，显存占用是推理阶段的2-3倍。报错中需要分配的[1800,1024,28,28] float张量单块就占用约5.4GB显存，该张量为RPN层生成的候选框特征，仅和模型配置、batch size相关，和训练数据集规模无关。
• 核查训练额外配置：确认是否开启了梯度全量记录、TensorBoard直方图存储、调试模式下的数值校验等额外功能，这类功能会额外占用30%以上的显存。同时确认TensorFlow Model Zoo的默认RPN配置，默认的候选框生成数量、非极大值抑制保留数量参数均为适配大规模数据集设置，小样本训练下不需要过高的参数值。

• 调整显存分配规则：在训练代码最开头加入以下配置，让TensorFlow按需分配显存而非预占全部显存：

import tensorflow as tf
gpus = tf.config.list_physical_devices('GPU')
if gpus:
  try:
    for gpu in gpus:
      tf.config.experimental.set_memory_growth(gpu, True)
  except RuntimeError as e:
    print(e)

• 继续下调batch size：在可用显存仅6GB的前提下，ResNet101 backbone的Faster R-CNN最高支持batch size为2，可先降到1跑通流程，再根据显存占用情况逐步上调。
• 开启混合精度训练：开启半精度训练可直接降低50%左右的显存占用，同时提升训练速度，在训练代码中加入以下配置即可：

from tensorflow.keras import mixed_precision
mixed_precision.set_global_policy('mixed_float16')

注意需要将优化器套入LossScaleOptimizer避免梯度下溢。

• 降低模型配置冗余项：
  • 将RPN层的first_stage_max_proposals参数从默认1000降到200-300区间
  • 将输入分辨率从640x640降到512x512，小样本训练下精度损失可以忽略
  • 替换轻量backbone，将ResNet101替换为ResNet50，小样本训练下精度差异极小
• 关闭冗余调试功能：训练阶段关闭TensorBoard的直方图、权重分布记录，关闭数值校验调试选项，减少非必要显存占用。
• 若仍报错可先使用CPU跑通训练流程，确认代码逻辑没有异常显存泄露后再切回GPU训练。

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