TensorFlow类实现模型内存占用过高问题求助

问题背景 #

原Functional API模型参数仅150万，切换到自定义keras.Model实现无监督训练后，出现构建计算图时内存占用飙升至200GB、训练启动耗时极长的异常，且训练中存在CPU/GPU负载下降但内存高位占用的情况。环境为M1 Pro（16GB内存+交换空间）、TensorFlow 2.13.0、Python 3.11.4，输入维度为(16, 387, 826, 1)。

核心问题定位 #

问题根源在于自定义Model类的实现不符合Keras内部机制规范，导致计算图构建冗余、内存管理失效，而非模型参数本身的规模问题。

针对性解决方案 #

1. 修复compile方法的错误实现 #

当前compile方法仅手动保存优化器和损失，未调用父类compile完成核心初始化，导致Keras无法正确配置梯度追踪、变量管理等模块，引发内存泄漏和图膨胀。

修正代码：

def compile(self, optimizer, loss, **kwargs):
    # 调用父类compile并传入必要参数，让Keras管理优化器和损失
    super(CustomFit, self).compile(optimizer=optimizer, loss=loss, **kwargs)

2. 规范子模型的注册方式 #

将Functional模型作为子组件嵌入时，需确保Keras能正确识别其参数并纳入管理，避免重复构建计算图。

优化__init__方法：

def __init__(self, model):
    super(CustomFit, self).__init__()
    # 将子模型包装为可训练层，确保参数被Keras正确追踪
    self.model = keras.layers.LayerWrapper(model)

3. 标准化train_step实现 #

无监督场景下的train_step必须严格遵循Keras张量操作规范，避免Python逻辑触发的冗余图构建。以下是标准模板，可根据你的损失逻辑调整：

def train_step(self, image):
    # 梯度追踪上下文，仅包裹前向传播
    with tf.GradientTape() as tape:
        # 显式设置training=True，确保BatchNorm等层处于训练模式
        predictions = self.model(image, training=True)
        # 无监督损失计算（示例为重构损失，替换为你的实际损失）
        loss = self.loss(image, predictions)
        # 累加正则化损失（如果模型包含Dropout、L2正则等）
        loss += sum(self.losses)
    
    # 计算梯度并更新参数
    gradients = tape.gradient(loss, self.trainable_variables)
    self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))
    
    # 返回标准化的指标字典，Keras依赖此更新训练状态
    return {"loss": loss}

4. M1 Pro环境专属优化 #

针对Apple Silicon平台，通过以下配置进一步降低内存占用并加速训练：

• 启用混合精度训练：

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

• 开启XLA加速（减少图构建时间和内存开销）：

  tf.config.optimizer.set_jit(True)
  

• 启用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)
  

5. 额外排查建议 #

• 确保数据加载使用tf.data.Dataset，避免纯Python迭代器引发的图追踪冗余；
• 临时开启 eager 执行调试：tf.config.run_functions_eagerly(True)，确认是否是计算图构建导致的内存问题（调试完成后关闭）。

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