使用Keras搭建图像去模糊GAN时Colab内存溢出与本地GPU OOM问题咨询

你基于DeblurGAN论文模型搭建的图像去模糊GAN训练代码，存在多处逻辑错误和资源使用不合理的问题，对应两个OOM问题的原因和修复方案如下：

一、代码存在的显性错误 #

• 日志与权重保存逻辑错位：当前write_logs、save_weights函数放在了epoch循环的外部，仅会在所有epoch跑完后执行一次。每个epoch产生的d_losses、gan_losses列表会持续堆积在内存中无法释放，是Colab内存溢出的核心原因之一。
• 变量重复赋值错误：epoch循环内出现x_train = dataset['sharp_img']的无效代码，你初始已经通过load_h5_dataset()加载了x_train和y_train，如果dataset是全局挂载的全量hdf5对象，该操作会重复把全量数据集加载到内存，进一步占用内存空间；如果dataset未定义还会直接触发运行时报错。
• 预测缓存未清理：每次调用g.predict都会在TensorFlow全局计算图中留下缓存，多轮迭代后缓存持续堆积，同时占用显存和内存空间。

二、Colab内存溢出修复方案 #

• 调整工具函数位置：将write_logs、save_weights缩进一层放到epoch循环内部，每个epoch跑完就计算损失平均值、写入日志，之后手动清空d_losses、gan_losses列表释放内存。
• 改用hdf5懒加载：不要用load_h5_dataset一次性把6G数据集全部加载到内存，改用h5py库的只读模式打开hdf5文件，训练时按batch索引读取数据，仅把当前需要的batch加载到内存即可，示例代码如下：

import h5py
with h5py.File('your_dataset_path.h5', 'r') as f:
    x_train = f['blur_img']
    y_train = f['sharp_img']
    # 训练循环内直接按索引取batch即可，不会全量加载数据集

• 关闭多余打印输出：移除不必要的print语句，减少内存占用。

三、本地显存溢出修复方案 #

本地报错为卷积层张量分配时显存不足，修复方式如下：

• 降低batch size：当前batch size设为16，叠加WGAN判别器多轮更新、VGG16感知损失中间张量的占用，6G显存无法支撑，建议先将batch size降到4或2测试。
• 替换预测调用方式：不要用g.predict生成伪图，直接用generated_images = g(image_blur_batch, training=False)，可避免predict产生的额外显存缓存。
• 启用显存动态分配：在代码开头加入如下配置，让TensorFlow按需分配显存而非一次性占满所有显存：

import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

• 开启混合精度训练：启用TensorFlow混合精度训练策略，用float16存储大部分张量，可减少近一半的显存占用：

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

• 优化感知损失计算：VGG16预训练模型仅加载到计算感知损失需要的中间层即可，同时设置VGG16所有层为不可训练，避免存储多余梯度状态占用显存。

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