Keras 3加载.keras模型反序列化错误：Dense层权重加载异常求助

迁移至Keras 3（TensorFlow v2.17.0、Keras v3.4.1）后，训练VAE模型时运行正常，但保存后加载出现权重加载错误。加载时提示z_mean、z_log_var等Dense层（或Conv2D层）预期变量未找到，尝试拆分保存架构和权重也无法解决，还出现优化器变量不匹配的警告。

核心原因 #

问题出在VAE的模型结构定义上：直接用函数式API拼接层而没有将整个VAE封装为自定义keras.Model子类，导致部分核心层（如z_mean、z_log_var）没有被模型实例正确跟踪，保存时未将这些层的权重写入文件，加载时自然找不到对应变量。

修复步骤 #

1. 将VAE封装为自定义Model子类 #

必须显式定义VAE类，包含encoder、decoder，并在call方法中实现前向传播，确保所有层都被模型跟踪。示例代码如下：

import keras
from keras import layers
import numpy as np
import keras.backend as K

latent_dim = 32
height, width, channels = 64, 64, 3  # 根据你的数据调整

# 注册采样函数（确保放在类定义前）
@keras.saving.register_keras_serializable()
def sampling(args):
    z_mean, z_log_var = args
    epsilon = K.random_normal(shape=(K.shape(z_mean)[0], latent_dim), mean=0., stddev=1.0)
    return z_mean + K.exp(z_log_var / 2) * epsilon

# 自定义VAE模型类
class VAE(keras.Model):
    def __init__(self, latent_dim, height, width, channels, **kwargs):
        super().__init__(**kwargs)
        self.latent_dim = latent_dim
        
        # 构建Encoder
        self.encoder_input = layers.Input(shape=(height, width, channels), name='encoder_input')
        x = layers.Conv2D(64, (3, 3), activation='relu', padding='same')(self.encoder_input)
        shape_before_flattening = K.int_shape(x)[1:]
        x = layers.Flatten()(x)
        self.z_mean = layers.Dense(latent_dim, name='z_mean')(x)
        self.z_log_var = layers.Dense(latent_dim, name='z_log_var')(x)
        self.z = layers.Lambda(sampling, output_shape=(latent_dim,), name='z')([self.z_mean, self.z_log_var])
        
        # 构建Decoder
        self.decoder_input = layers.Input(K.int_shape(self.z)[1:])
        x = layers.Dense(np.prod(shape_before_flattening))(self.decoder_input)
        x = layers.Reshape(shape_before_flattening)(x)
        self.decoder_output = layers.Conv2D(channels, (3, 3), activation='sigmoid', padding='same')(x)
        
        # 定义完整模型的输入输出
        self.encoder = keras.Model(self.encoder_input, [self.z_mean, self.z_log_var, self.z], name='encoder')
        self.decoder = keras.Model(self.decoder_input, self.decoder_output, name='decoder')
        
    def call(self, inputs):
        z_mean, z_log_var, z = self.encoder(inputs)
        reconstruction = self.decoder(z)
        # 添加VAE自定义损失（重构损失+KL散度）
        reconstruction_loss = keras.losses.binary_crossentropy(K.flatten(inputs), K.flatten(reconstruction))
        reconstruction_loss *= height * width * channels
        kl_loss = -0.5 * K.sum(1 + z_log_var - K.square(z_mean) - K.exp(z_log_var), axis=-1)
        total_loss = K.mean(reconstruction_loss + kl_loss)
        self.add_loss(total_loss)
        return reconstruction

2. 训练与保存模型 #

使用自定义类实例化模型，训练后直接保存：

# 实例化模型
vae = VAE(latent_dim, height, width, channels)
vae.compile(optimizer='adam')

# 训练（假设x_train是你的训练数据）
vae.fit(x_train, epochs=10, batch_size=32)

# 保存模型
vae.save('vae_model.keras')

3. 加载模型 #

加载时传入自定义对象（包含采样函数和VAE类）：

custom_objects = {'sampling': sampling, 'VAE': VAE}
loaded_vae = keras.models.load_model('vae_model.keras', custom_objects=custom_objects)

为什么之前的方法失败？ #

• 直接用函数式API拼接层时，最终的模型可能只以encoder_input到decoder_output为输入输出，没有显式包含z_mean、z_log_var这些中间层，导致保存时这些层的权重被忽略。
• 拆分保存json和weights的方式在Keras3中已不推荐，且同样会因为模型结构未完整定义而丢失层权重。

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