训练时通过LambdaCallback创建TensorFlow Keras子模型遇图断开错误

核心原因：训练上下文的计算图隔离 #

在Keras训练过程中，框架会自动构建专属训练计算图，这个图包含梯度计算、优化器更新、正则化等训练相关操作逻辑。此时原模型中的张量（比如通过model.get_layer("some_name").input获取的张量）会绑定在这个训练图内。

当你在回调中尝试创建新模型时，新模型默认会在当前训练图上下文里构建，但训练图的结构是为原模型的完整训练流程设计的，部分张量链路被训练操作占用或隔离，导致新模型无法正确追踪从输入到输出的完整数据流，最终触发“图断开”错误。

你提到的几种正常场景，本质是脱离了训练图的约束：

• 训练上下文外：此时使用默认计算图，原模型的张量未被训练操作绑定，新模型可正常构建数据流链路。
• 克隆模型后：克隆操作会生成与原模型权重共享但计算图独立的实例，其内部张量属于新计算图，和原训练图无关联，因此能正常构建子模型。

可行解决方案 #

方案1：在回调中先克隆原模型，再构建子模型 #

克隆操作会生成独立的模型实例，避开训练图约束：

from tensorflow.keras.models import clone_model
import tensorflow as tf

def create_submodel(original_model):
    # 克隆原模型并同步权重
    cloned_model = clone_model(original_model)
    cloned_model.set_weights(original_model.get_weights())
    # 从克隆模型中获取层构建子模型
    input_layer = cloned_model.get_layer("some_name").input
    other_input_layer = cloned_model.get_layer("other_name").input
    outputs = cloned_model.get_layer("target_output_layer").output
    submodel = tf.keras.Model(inputs=[input_layer, other_input_layer], outputs=outputs)
    return submodel

# 在LambdaCallback中使用
epoch_end_callback = tf.keras.callbacks.LambdaCallback(
    on_epoch_end=lambda epoch, logs: create_submodel(model).save(f"submodel_epoch_{epoch}.h5")
)

方案2：提前定义子模型结构，训练时同步权重 #

如果子模型结构固定，可提前定义并利用层名称匹配共享原模型权重：

import tensorflow as tf

# 提前定义子模型（层名称与原模型对应层完全一致）
def build_submodel():
    input_layer = tf.keras.Input(shape=(...), name="some_name")
    other_input_layer = tf.keras.Input(shape=(...), name="other_name")
    # 复刻原模型中间层结构，确保层名称匹配
    x = tf.keras.layers.Dense(64, name="dense_1")(input_layer)
    y = tf.keras.layers.Dense(32, name="dense_2")(other_input_layer)
    combined = tf.keras.layers.Concatenate(name="concat")([x, y])
    outputs = tf.keras.layers.Dense(10, name="target_output")(combined)
    return tf.keras.Model(inputs=[input_layer, other_input_layer], outputs=outputs)

submodel = build_submodel()
# 初始同步原模型权重
submodel.set_weights(model.get_weights())

# 回调中直接保存子模型，训练时权重会自动同步
epoch_end_callback = tf.keras.callbacks.LambdaCallback(
    on_epoch_end=lambda epoch, logs: submodel.save(f"submodel_epoch_{epoch}.h5")
)

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