Keras构建模型时GPU内存占满引发训练OOM问题求助

可能原因 #

• GPU内存增长设置时机错误：tf.config.experimental.set_memory_growth必须在所有TensorFlow/Keras相关操作（比如导入模块、加载数据、定义模型层）之前执行，否则内存预分配已经完成，设置不会生效。你可能是在执行了部分TensorFlow操作后才运行配置代码，导致规则没被应用。
• 模型结构的内存预分配：U-net包含大量卷积层与跳跃连接，TensorFlow在构建模型计算图时，会提前为模型参数、中间张量预分配GPU内存，尤其是TF2.10仍保留部分静态图兼容逻辑，这会导致模型构建完成后就占用大量内存。
• 残留进程未释放：如果之前的代码异常终止或未正确清理资源（比如未销毁的张量、会话），GPU内存会被持续占用，只能通过重启内核或系统释放。

解决方法 #

1. 调整GPU配置的执行顺序 #

把GPU内存增长的配置代码放在所有TensorFlow/Keras导入和操作的最开头，确保在任何初始化动作前生效：

import tensorflow as tf

# 优先配置GPU
gpus = tf.config.list_physical_devices('GPU')
if gpus:
    try:
        tf.config.experimental.set_memory_growth(gpus[0], True)
        print("GPU内存增长模式已启用")
    except RuntimeError as e:
        print(e)

# 之后再导入Keras模块、加载数据、构建模型
from tensorflow.keras import layers, models
# ... 后续代码

2. 直接限制GPU内存占用比例 #

如果内存增长模式无效，可以直接指定TensorFlow能使用的GPU内存上限，比如限制为8GB显卡的70%（约5734MiB）：

import tensorflow as tf

gpus = tf.config.list_physical_devices('GPU')
if gpus:
    try:
        tf.config.set_logical_device_configuration(
            gpus[0],
            [tf.config.LogicalDeviceConfiguration(memory_limit=5734)]
        )
        logical_gpus = tf.config.list_logical_devices('GPU')
        print(f"{len(gpus)} 物理GPU, {len(logical_gpus)} 逻辑GPU")
    except RuntimeError as e:
        print(e)

3. 主动清理GPU残留资源 #

在代码开头和训练结束后添加资源清理操作，避免残留张量占用内存：

import gc
import tensorflow as tf

# 运行前清理
gc.collect()
tf.keras.backend.clear_session()

# ... 构建模型、训练代码 ...

# 训练结束后清理
tf.keras.backend.clear_session()
gc.collect()

如果使用Jupyter Notebook，每次运行训练代码前建议重启内核，彻底清除之前的残留进程。

4. 启用混合精度减少内存占用 #

通过混合精度训练，将部分张量从float32转为float16，大幅降低内存占用，同时保证关键层的精度：

import tensorflow as tf
from tensorflow.keras import mixed_precision, layers, models

# 启用混合精度
mixed_precision.set_global_policy('mixed_float16')

# 构建U-net模型，注意输出层指定float32避免精度损失
def build_unet():
    # ... 编码器、解码器结构 ...
    # 输出层设置dtype='float32'
    output = layers.Conv2D(1, (1,1), activation='sigmoid', dtype='float32')(last_layer)
    model = models.Model(inputs=input_layer, outputs=output)
    return model

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