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