如何正确保存与加载嵌套自定义 Keras 模型（避免层命名冲突与权重加载失败）-Python教程-PHP中文网

Keras 自定义模型（尤其是含子模型嵌套结构的）在使用 model.save() 和 keras.saving.load_model() 时，若在 train_step 中绕过主模型 call() 直接调用子模型方法，会导致层命名不一致、变量注册缺失，从而引发“Shape mismatch”或“expected 0 variables but received 2”等反序列化错误。

keras 自定义模型（尤其是含子模型嵌套结构的）在使用 `model.save()` 和 `keras.saving.load_model()` 时，若在 `train_step` 中绕过主模型 `call()` 直接调用子模型方法，会导致层命名不一致、变量注册缺失，从而引发“shape mismatch”或“expected 0 variables but received 2”等反序列化错误。

在 TensorFlow/Keras 2.15+ 中，自定义 keras.Model 子类的可保存性不仅依赖于 get_config() 和 from_config() 的正确实现，更关键的是模型图结构的完整性与层生命周期的一致性。你遇到的 ValueError: Cannot assign value to variable 'dense_4/kernel:0': Shape mismatch 并非真实形状变化，而是加载时 Keras 尝试将权重映射到未被正确定义或未被 call() 触发构建的层实例上——这些层在保存时未被完整注册进模型的内部变量管理器（self._trainable_weights / self._non_trainable_weights），导致加载时权重名称与目标层不匹配。

? 根本原因：绕过 call() 破坏了层构建与变量注册链

观察你的 HighNet.train_step：

z = self.encoder(data)           # ✅ 正确：触发 encoder.call() → 完整构建 + 注册所有层变量
noise_z = self.generator(noise_tf)  # ✅ 同理

这本身是安全的 —— 但问题往往隐藏在更深层：若 Encoder 或 Generator 类自身未在 __init__ 中完成所有层的显式构建（即未调用 build() 或未在 call() 中首次调用其子层），其内部 Dense 层可能处于“延迟构建”状态。而当 train_step 中直接调用 self.encoder(data) 时，Keras 会动态构建该子模型并注册变量；但保存时仅序列化了当时已注册的变量集合。若后续训练/调用路径导致不同层被激活（如 BatchNormalization 的 training=True 分支），或子模型在 __init__ 中未统一初始化（如 build() 被手动调用但未绑定到实例属性），就会造成保存/加载时变量视图不一致。

更隐蔽的问题来自你的最小复现示例：

# ❌ 危险：在 train_step 中绕过主模型 call()
y_pred = self.net2(self.net1(x))  # → net1/net2 的 call() 被调用，但其层未被 HighNet 的 graph context 完整捕获

# ✅ 安全：所有子模型调用必须经由主模型的 call() 链路
def call(self, inputs):
    z = self.net1(inputs)
    x = self.net2(z)
    return z, x

# 在 train_step 中：
z, y_pred = self(x)  # → 触发完整 call 图，确保所有层被纳入模型拓扑

Keras 的 SavedModel/Keras 格式依赖于静态计算图的可追踪性。直接在 train_step 内部调用子模型方法（即使它们是 keras.Model）会绕过主模型的 __call__ 调度机制，导致：

子模型层未被加入主模型的 layers 列表；
其 trainable_weights 不出现在 self.trainable_weights 中（尽管实际参与梯度更新）；
序列化时仅保存了“可见”的权重，而加载时 Keras 按完整 call() 图重建层，发现权重数量/名称不匹配。

✅ 正确实践：三层保障策略

1. 子模型必须通过 call() 统一入口接入

重写 HighNet.call()，明确暴露所有子模型逻辑：

@tf.keras.saving.register_keras_serializable()
class HighNet(keras.Model):
    def __init__(self, encoder, decoder, args, **kwargs):
        super().__init__(**kwargs)
        self.encoder = encoder
        self.decoder = decoder
        self.args = args
        self.generator = Generator(args).build()  # ✅ 确保 build() 返回 Model 实例并赋值
        self.reconstruction_loss_tracker = keras.metrics.Mean("cw_reconstruction_loss")

    def call(self, inputs, training=None):
        """✅ 强制所有前向传播经过此入口"""
        z = self.encoder(inputs, training=training)
        return self.decoder(z, training=training)

    def generate_latent(self, noise, training=None):
        """辅助方法：仍需通过 generator.call()"""
        return self.generator(noise, training=training)

    def train_step(self, data):
        with tf.GradientTape() as tape:
            # ✅ 全部调用走 call() 链路
            z_real = self.encoder(data, training=True)
            batch_size = tf.shape(z_real)[0]
            noise_np = np.random.normal(0, 1, size=(batch_size, self.args["noise_dim"]))
            noise_tf = tf.convert_to_tensor(noise_np, dtype=tf.float32)
            z_fake = self.generate_latent(noise_tf, training=True)
            cw_reconstruction_loss = tf.math.log(
                cw_sampling_silverman(z_real, z_fake)
            )
        # ... 梯度更新保持不变

2. 确保子模型 build() 可重复且幂等

修改 Encoder.build() 等方法，避免返回新 Model 实例（这会切断权重引用）：

class Encoder(keras.Model):
    def __init__(self, args, **kwargs):
        super().__init__(**kwargs)
        self.args = args
        self.flatten = layers.Flatten()
        self.dense1 = layers.Dense(256, name="dense1")  # ✅ 显式命名
        self.dense2 = layers.Dense(args["latent_dim"], name="z")
        self.activation = layers.Activation("relu")

    def call(self, inputs, training=None):  # ✅ 移除 build()，用 call() 构建
        x = self.flatten(inputs)
        x = self.dense1(x)
        x = self.activation(x)
        z = self.dense2(x)
        return z

⚠️ 注意：build() 方法在 keras.Model 子类中不应返回新 Model 实例，而应仅做参数初始化。你原代码中 Encoder.build() 返回 keras.Model(...)，这创建了一个与 self 无关的独立模型，其权重不会被 HighNet 管理 —— 这正是 dense_4/kernel:0 名称混乱的根源（Keras 自动生成编号）。

3. 保存前验证模型图完整性

在 model.save() 前插入检查：

def validate_model_structure(model):
    print("=== Model Weight Summary ===")
    for i, w in enumerate(model.trainable_weights):
        print(f"Weight {i}: {w.name} → shape {w.shape}")
    print(f"Total trainable weights: {len(model.trainable_weights)}")
    model.summary()  # ✅ 强制触发完整图构建

# 使用示例
validate_model_structure(model)
model.save("high_model.keras", save_format="keras")

? 关键注意事项

禁用 pickle/dill：Keras 模型包含 C++ 后端对象（如 tf.Variable），无法被纯 Python 序列化器安全处理。
get_config() 必须可序列化：self.encoder 是 keras.Model 实例，不能直接存入 config（会报 TypeError: not JSON serializable）。应改为存储构造参数：
```
def get_config(self):
    return {
        "args": self.args,
        "latent_dim": self.args["latent_dim"],
        "noise_dim": self.args["noise_dim"],
    }
```
并在 from_config 中重建子模型（需确保子模型也支持序列化）。
TensorFlow 版本兼容性：TF 2.15 对 Keras 2.15 的序列化有严格要求，务必确认所有自定义层/模型均添加 @keras.saving.register_keras_serializable()，且无闭包引用外部不可序列化对象（如 numpy.ndarray、文件句柄）。

✅ 最终验证流程

# 1. 训练后立即验证
model.save("high_model.keras")
print("✅ Save succeeded")

# 2. 加载并测试前向传播
loaded = keras.saving.load_model("high_model.keras")
test_input = tf.random.normal((1, 28, 28, 1))
_ = loaded(test_input)  # ✅ 必须能执行 call()

# 3. 检查权重一致性
assert len(model.trainable_weights) == len(loaded.trainable_weights)
for w_orig, w_load in zip(model.trainable_weights, loaded.trainable_weights):
    assert w_orig.shape == w_load.shape, f"Shape mismatch: {w_orig.name} vs {w_load.name}"

遵循以上原则，即可彻底解决因嵌套模型调用方式不当导致的保存/加载失败问题，确保模型在训练、部署、协作场景中具备完全可复现性。