问题重现 #

代码在CPU环境下可正常完成训练，但GPU环境中RAM持续增长，首个训练步骤完成后内存未释放，最终触发内存不足错误。

核心原因分析 #

• 禁用@tf.function导致Eager模式内存失控
  代码注释掉了@tf.function装饰器，强制在Eager模式下运行训练步骤。Eager模式会为每个操作创建独立的张量对象，无法利用TensorFlow的图优化和内存复用机制，大量临时张量堆积在GPU内存中无法及时回收。

• 循环计算Loss产生冗余张量
  手动循环每个时间步累加Loss，每一次迭代都会生成新的中间张量，这些张量会被梯度磁带全程追踪，占用额外内存。

• Metric未重置导致内存累积
  self.metric在训练过程中从未重置，每次step都会将新的计算结果追加到历史数据中，随着训练轮次增加，内存占用持续上升。

• 手动内存回收无效
  TensorFlow的GPU内存由内部内存池管理，手动执行del和gc.collect()无法直接释放GPU显存，反而可能干扰TensorFlow的内存优化逻辑。

针对性修复方案 #

1. 恢复@tf.function装饰器

重新启用图模式，让TensorFlow自动优化内存使用和操作执行：

@tf.function
def train_step(self, data):
    # 原训练逻辑

2. 向量运算替代Loss循环计算

用向量化操作替换逐时间步的循环，减少中间张量生成：

# 替换原Loss循环计算
loss = tf.reduce_mean(tf.square(masked_y_true - masked_outputs))
# 若需按时间步平均，可改为：
loss = tf.reduce_mean(tf.square(masked_y_true - masked_outputs), axis=[0,1,2])
loss = tf.reduce_mean(loss)

3. 每个Step后重置Metric

在train_step末尾添加重置逻辑，避免历史数据累积：

result = {"loss": loss, "mae": self.metric.result()}
self.metric.reset_state()  # 重置指标
return result

4. 移除无效的手动内存回收代码

删除del outputs_ta, ...和gc.collect()语句，交给TensorFlow内存管理器处理显存释放。

5. 优化TensorArray使用（图模式下）

在图模式中，为TensorArray添加clear_after_read=True参数，减少内存占用：

outputs_ta = tf.TensorArray(dtype=tf.float32, size=seq_len, dynamic_size=False, clear_after_read=True)

额外优化建议 #

• 检查self.model中的LSTM层是否存在状态泄漏，确保每次训练步的初始状态正确重置。
• 启用TensorFlow的内存增长模式，避免一次性占用全部GPU显存：

gpus = tf.config.list_physical_devices('GPU')
if gpus:
    try:
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
    except RuntimeError as e:
        print(e)

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