TensorFlow-Keras加载权重后计算图体积大幅增大问题咨询

这种每层新增约10个操作、内存大幅上升的情况绝对不是预期行为，在TensorFlow 1.14 + Keras 2.3.0的环境里，这大概率是静态计算图模式下加载权重时的机制缺陷导致的额外计算节点残留。

原因分析 #

TF1.x采用静态计算图模式，model.load_weights() 在加载仅权重的文件时，会为每个模型变量生成对应的赋值操作（比如Assign类节点）来完成权重写入。如果加载前计算图已有残留节点，或者权重匹配过程中产生了临时操作，这些节点都会被保留在计算图中，不会自动清理。累积下来，每层就会新增不少额外操作，直接推高内存占用。

解决方案 #

• 重置计算图后再加载权重
  在加载权重前，先重置默认计算图，再重新构建模型，这样能彻底避免旧图的节点残留。示例代码：

  import tensorflow as tf
  from tensorflow.keras.models import Model
  
  # 重置默认计算图，清空所有残留节点
  tf.reset_default_graph()
  
  # 重新定义你的自定义模型
  def build_your_model():
      inputs = tf.keras.Input(shape=(...))
      x = tf.keras.layers.Dense(512, activation='relu')(inputs)
      # 这里替换成你的完整模型结构
      outputs = tf.keras.layers.Dense(10)(x)
      return Model(inputs=inputs, outputs=outputs)
  
  model = build_your_model()
  # 加载权重
  model.load_weights(path_to_existing_weights)
  

• 先加载模型结构再加载权重
  如果可以提前保存模型结构，先加载结构再加载权重的方式会更稳定，减少权重匹配时的额外操作：

  # 先保存模型结构（训练时执行一次即可）
  model_json = model.to_json()
  with open("model_structure.json", "w") as f:
      f.write(model_json)
  
  # 加载时的操作
  tf.reset_default_graph()
  from tensorflow.keras.models import model_from_json
  
  with open("model_structure.json", "r") as f:
      model_json = f.read()
  model = model_from_json(model_json)
  model.load_weights(path_to_existing_weights)
  

• 手动赋值权重，精准控制节点生成
  跳过Keras的自动加载逻辑，手动遍历权重文件和模型变量完成赋值，避免自动生成额外操作：

  from tensorflow.python.training.checkpoint_utils import get_variable_names, load_variable
  
  # 从权重文件中读取所有变量及其值
  checkpoint_vars = {
      name: load_variable("./your_checkpoint_dir", name) 
      for name in get_variable_names("./your_checkpoint_dir")
  }
  
  # 遍历模型变量，手动完成赋值
  for var in model.variables:
      # 去掉变量名后的":0"后缀，匹配权重文件中的命名
      var_name = var.name.split(":")[0]
      if var_name in checkpoint_vars:
          var.assign(checkpoint_vars[var_name])
  

• 升级到TensorFlow 2.x（若业务允许）
  TF1.x的静态计算图本身就容易出现节点累积问题，升级到TF2.x后默认采用动态图模式，计算节点会按需创建和销毁，这类内存暴涨的情况基本不会再出现。当然这需要你调整代码适配TF2.x的语法，比如关闭兼容模式、调整层的定义方式等。

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