调用Keras model.fit()后出现内存泄漏问题求助

问题描述 #

我编写了名为Material的自定义Keras Layer并在模型中使用，但每次调用model.fit()（指函数调用，而非每个epoch）后，内存占用都会不断增加，到约5GB时增速放缓但仍持续增长，该问题在CPU和GPU环境下均存在。想请教这是什么原因？自定义Layer是否存在问题？

复现代码 #

import tensorflow as tf
from tensorflow import keras
import numpy as np

class Material(keras.layers.Layer):
    def __init__(self):
        super().__init__()
        self.table = tf.Variable(
            initial_value=np.ones(shape=(6, 8, 8), dtype="float32"),
            trainable=True,
        )
        self.k = tf.Variable(initial_value=8., trainable=True, dtype='float32')
    
    def call(self, inp):
        material_white = tf.reduce_sum(inp[..., 0] * self.table, axis=(-1, -2, -3))
        material_black = tf.reduce_sum(inp[..., 1] * self.table, axis=(-1, -2, -3))
        material_white = tf.maximum(material_white, .01)
        material_black = tf.maximum(material_black, .01)
        return tf.math.log(material_white / material_black) * self.k

def get_model() -> keras.Model:
    inp = keras.Input((6, 8, 8, 2),)
    material = Material()
    out = keras.activations.sigmoid(material(inp))
    return keras.Model(inputs=inp, outputs=out)

model = get_model()
model.compile(optimizer=keras.optimizers.Adam(learning_rate=.01), loss='mean_squared_error')

board_tables = np.random.sample((100000, 6, 8, 8, 2),)
outcome = np.random.sample((100000,),)

while 1:
    model.fit(board_tables, outcome, batch_size=200, epochs=2)

原因分析 #

• 自定义Layer本身无内存泄漏问题：你的Material Layer逻辑中，变量初始化和前向计算都是标准操作，没有创建无限增长的变量或保留不必要的张量引用，Layer本身不存在内存泄漏点。
• 循环调用model.fit()的日志与计算图累积：每次调用model.fit()时，TensorFlow默认会记录训练过程的日志信息（如损失值、指标），同时在动态图模式下，每次调用可能生成新的计算图片段并保留，导致内存中累积大量冗余的计算图节点和日志数据，这是内存持续增长的核心原因。

解决建议 #

• 合并训练轮次，减少model.fit()调用次数：把while 1循环改成一次性设置足够多的epochs，避免反复调用fit生成冗余日志和计算图：

  # 直接设置大数量epochs，替代无限循环
  model.fit(board_tables, outcome, batch_size=200, epochs=1000)
  

• 关闭不必要的日志输出：如果必须保留循环调用，关闭默认日志打印，减少日志数据累积：

  while True:
      model.fit(board_tables, outcome, batch_size=200, epochs=2, verbose=0)
      # 手动打印必要信息，替代默认日志
      print(f"当前最后一轮损失: {model.history.history['loss'][-1]}")
  

• 用tf.function装饰call方法：将前向计算转为静态图模式，避免每次调用生成新的计算图片段：

  class Material(keras.layers.Layer):
      def __init__(self):
          super().__init__()
          self.table = tf.Variable(
              initial_value=np.ones(shape=(6, 8, 8), dtype="float32"),
              trainable=True,
          )
          self.k = tf.Variable(initial_value=8., trainable=True, dtype='float32')
      
      @tf.function  # 添加静态图装饰器
      def call(self, inp):
          material_white = tf.reduce_sum(inp[..., 0] * self.table, axis=(-1, -2, -3))
          material_black = tf.reduce_sum(inp[..., 1] * self.table, axis=(-1, -2, -3))
          material_white = tf.maximum(material_white, .01)
          material_black = tf.maximum(material_black, .01)
          return tf.math.log(material_white / material_black) * self.k
  

• 手动清除训练历史：每次训练后删除历史记录对象，避免日志数据占用内存：

  while True:
      history = model.fit(board_tables, outcome, batch_size=200, epochs=2)
      del history  # 清除历史记录释放内存
  

内容的提问来源于stack exchange，提问作者Усердный бобёр