在Google Colab中Keras的clear_session()无法正常工作的问题

在Colab中重复运行Keras模型代码时，即便调用K.clear_session()，仍可能出现内存耗尽的情况，核心原因是Python对模型对象的引用未被释放，加上TensorFlow Eager模式下的资源管理特性，导致旧模型占用的内存无法被回收。以下是几种有效的解决方案：

1. 先清理模型引用再执行会话清理 #

K.clear_session()仅清理Keras内部会话，但如果模型对象仍被Python变量引用，垃圾回收器无法释放其内存。需按顺序执行以下步骤：

# 模型训练完成后
m.fit(X,y)
# 删除模型变量引用
del m
# 触发Python垃圾回收
import gc
gc.collect()
# 清理Keras会话
K.clear_session()

2. 配合重置TensorFlow默认计算图 #

部分场景下，仅清理Keras会话不足以释放所有图资源，需额外重置默认计算图：

import tensorflow as tf
# ...模型训练代码...
del m
gc.collect()
K.clear_session()
tf.compat.v1.reset_default_graph()

3. 将训练逻辑封装为函数 #

把模型定义、训练代码放入函数中，函数执行完毕后局部变量会自动回收，避免全局变量导致的内存泄漏：

import numpy as np
from tensorflow.keras import Sequential
from tensorflow.keras.layers import Dense
from tensorflow.keras import backend as K
import gc
import tensorflow as tf

X = np.zeros([10, 10000])
y = np.zeros([10, 10000])

def train_model(X, y):
    m = Sequential([Dense(10000, input_shape=(10000,)), Dense(10000), Dense(10000), Dense(10000)])
    m.compile(loss='mse')
    m.summary()
    m.fit(X,y)
    return m

# 每次运行训练流程
model = train_model(X, y)
del model
gc.collect()
K.clear_session()
tf.compat.v1.reset_default_graph()

4. 优化模型与训练策略 #

示例中的模型参数规模极大（每层10000神经元），本身就会占用大量内存，可通过以下方式降低内存压力：

• 减少神经元数量，缩小模型规模
• 启用混合精度训练：tf.keras.mixed_precision.set_global_policy('mixed_float16')
• 改用更小的batch size，避免一次性加载全量数据训练

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