使用Keras fit()训练CycleGAN时内存持续增长问题排查

• 检查Dataset迭代器的重复逻辑：如果同时在dataset.repeat()和fit的epochs参数中设置循环，会导致迭代器嵌套重复，引发内存累积。确保只在一处控制循环次数——要么Dataset层面仅做必要重复，fit的epochs设为1；要么移除Dataset的repeat()，由fit负责epoch循环。

• 调整Dataset预取策略：暂时去掉dataset.prefetch(tf.data.AUTOTUNE)，替换为固定预取数量（比如prefetch(2)）。AUTOTUNE在部分场景下会导致预取队列累积未释放的张量，尤其是fit循环重复迭代数据集时。

• 修改缓存存储方式：如果当前用的是默认内存缓存cache()，改成磁盘缓存cache('/tmp/dataset_cache')，强制将缓存数据写入磁盘，避免内存中堆积重复引用的张量。

• 清理回调函数的冗余存储：检查所有训练回调，比如ModelCheckpoint是否保存了过多旧模型文件，TensorBoard是否无限制记录历史数据。可以设置ModelCheckpoint仅保存最优模型，给TensorBoard添加max_queue参数限制日志队列大小，减少内存占用。

• 手动实现epoch循环：放弃使用fit()，手动编写训练循环，每个epoch结束后显式释放资源：

import gc
for epoch in range(EPOCHS):
    for batch in dataset:
        train_step(batch)
    tf.keras.backend.clear_session()
    gc.collect()

这种方式能直接控制资源释放，避免fit内部隐式的资源累积。

• 检查数据处理的张量追踪：确保自定义生成器和Dataset映射函数中，所有操作都用TF算子实现，避免混合Python原生变量。未被正确追踪的张量会滞留在计算图中无法释放，引发内存泄漏。

• 固定数据集元素形状：确认每个batch的张量形状、dtype完全一致。如果存在动态尺寸（比如不同batch图片大小不一），TF会不断创建新计算图分支，导致内存持续增长。可以在映射函数中强制固定尺寸，比如tf.image.resize(image, [256,256])。

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