Google Colab训练CNN出现OOM错误如何有效释放GPU显存？

问题背景 #

在Google Colab环境中循环训练CNN模型，使用tf.datasets提供的eurosat/rgb/数据集，模型与数据集规模均不大，但训练时触发ResourceExhaustedError显存耗尽错误。已尝试在每轮训练结束后调用以下清理函数，此前可正常运行，当前最小模型也会触发报错，报错时nvidia-smi显示Tesla K80的11G显存已占用11G左右，几乎被占满。

原有清理代码：

def reset_tensorflow_keras_backend():
    # 仍需进一步验证，目前看来足够使用
    import tensorflow as tf
    import tensorflow.keras as keras
    tf.keras.backend.clear_session()
    tf.compat.v1.reset_default_graph()
    _ = gc.collect()

有效解决方法 #

• 优化原有清理函数，新增显式显存释放逻辑
  在导入TensorFlow后第一时间开启显存动态分配，避免TensorFlow一次性抢占全部显存，每轮训练结束后调用优化后的清理函数：

import gc
import tensorflow as tf
from tensorflow.keras import backend as K

# 初始化阶段就开启显存动态分配
if tf.config.list_physical_devices('GPU'):
    for gpu in tf.config.list_physical_devices('GPU'):
        tf.config.experimental.set_memory_growth(gpu, True)

def reset_tensorflow_keras_backend():
    # 清理Keras会话
    K.clear_session()
    # 重置默认计算图
    tf.compat.v1.reset_default_graph()
    # 重置GPU显存统计，触发显存碎片整理
    tf.config.experimental.reset_memory_stats('GPU:0')
    # 触发全量垃圾回收
    gc.collect()

• 调整数据集加载策略
  使用tf.data加载数据时，不要用.cache()直接把全量数据缓存到显存，改为缓存到本地磁盘降低显存占用：

# 原缓存到显存的写法（显存占用高）
# dataset = dataset.cache().shuffle(1000).batch(batch_size).prefetch(tf.data.AUTOTUNE)
# 修改为缓存到磁盘
dataset = dataset.cache(filename="./eurosat_cache").shuffle(1000).batch(batch_size).prefetch(tf.data.AUTOTUNE)

• 降低训练批次大小
  可将原有batch_size减半测试，Colab分配的Tesla K80本身显存较小，运行时间久了会有残留进程占用部分显存，降低batch_size可直接降低训练时的显存峰值。
• 显式清空全部GPU显存
  如果以上方法都无效，直接杀掉当前GPU上的所有计算进程，运行后等待10秒Colab会自动重连，显存会被完全清空，无需重启整个实例：

!fuser -v /dev/nvidia* | awk '{for(i=1;i<=NF;i++)if($i~/^[0-9]+$/)print $i}' | xargs kill -9 2>/dev/null

• 模型结构优化
  报错指向全连接层的正则化计算，可将模型最后的大尺寸全连接层替换为全局平均池化层，可降低70%以上的全连接层显存占用。

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