TensorFlow中每轮训练后GPU内存大幅增长是否属正常行为？

这种每轮训练新增GPU内存占用的情况不属于正常行为，大概率是代码存在内存泄漏，而非TensorFlow默认需要存储大量历史数据。以下是常见原因和解决方向：

• 计算图重复构建：如果把模型定义、损失函数初始化这类操作放在训练循环内部，每轮都会生成新的计算图节点，导致内存持续累积。要把模型、优化器、损失函数的创建逻辑都放在训练循环外。
• 临时张量未回收：训练中手动创建的临时张量、自定义变量，或者tf.GradientTape的引用没有及时释放，都会留存内存。每轮训练结束后可以调用tf.keras.backend.clear_session()清理会话；使用GradientTape时，要确保它的作用域只在单轮训练内，不要在循环外保留磁带对象。
• 回调函数冗余存储：自定义回调如果保存了每轮的权重、中间层输出这类大张量，或者ModelCheckpoint设置为保存所有轮次权重而不覆盖，会快速占用内存。检查回调逻辑，只保留必要数据——比如只保存最优权重，不要存储每轮的中间结果。
• 数据集重复加载：如果在训练循环内重复创建数据集对象，每轮都会生成新的数据张量。确保数据集是一次性构建的，用cache()或prefetch()优化加载流程，避免循环内重复初始化数据集。

关于训练时需保留的数据：
TensorFlow默认只会保留当前轮次的模型权重、优化器状态（比如动量项）、当前批次的输入输出张量，以及损失值的累加器。不会自动存储历史权重，除非你手动通过代码或回调实现了这个功能。

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