Keras 2.10.0下自定义VAE模型无法保存的技术求助

问题背景 #

环境：Keras 2.10.0
自定义Variational Autoencoder（VAE）训练完成后无法保存，报错信息如下：

ValueError: Model <main.VAE object at 0x00000278F0A96E60> cannot be saved either because the input shape is not available or because the forward pass of the model is not defined.To define a forward pass, please override Model.call(). To specify an input shape, either call build(input_shape) directly, or call the model on actual data using Model(), Model.fit(), or Model.predict(). If you have a custom training step, please make sure to invoke the forward pass in train step through Model.__call__, i.e. model(inputs), as opposed to model.call().

报错原因 #

自定义VAE类继承自keras.Model但未实现call()方法，Keras保存模型时需要明确的前向传播逻辑来构建完整计算图；同时训练时直接调用encoder和decoder的方式，没有通过模型自身的前向传播入口，导致Keras无法追踪模型的输入输出结构。

解决方案 #

1. 完善VAE类的核心逻辑 #

为VAE类添加call()方法定义前向传播，同时修正get_config和添加from_config方法，确保子模型能正确序列化：

class VAE(keras.Model):
    def __init__(self, encoder, decoder, **kwargs):
        super().__init__(**kwargs)
        self.encoder = encoder
        self.decoder = decoder
        self.total_loss_tracker = keras.metrics.Mean(name="total_loss")

    @property
    def metrics(self):
        return [self.total_loss_tracker]

    # 新增call方法，定义完整前向传播路径
    def call(self, inputs):
        z_mean, z_log_var, z = self.encoder(inputs)
        reconstruction_cont, reconstruction_class = self.decoder(z)
        return reconstruction_cont, reconstruction_class, z_mean, z_log_var

    def train_step(self, data):
        with tf.GradientTape() as tape:
            # 改为调用模型自身的前向方法，统一逻辑
            reconstruction_cont, reconstruction_class, z_mean, z_log_var = self(data)

            data_cont = data[:, :4]
            data_class = data[:, 4:]

            reconstruction_loss_cont = keras.losses.mean_squared_error(data_cont, reconstruction_cont)
            reconstruction_loss_class = keras.losses.sparse_categorical_crossentropy(data_class, reconstruction_class)
            kl_loss = -0.5 * (1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var))
            kl_loss = tf.reduce_mean(kl_loss, axis=1)

            total_loss = 0.8*reconstruction_loss_cont + 0.2*reconstruction_loss_class + kl_loss

        grads = tape.gradient(total_loss, self.trainable_weights)
        self.optimizer.apply_gradients(zip(grads, self.trainable_weights))
        self.total_loss_tracker.update_state(total_loss)

        return {"loss": self.total_loss_tracker.result()}

    def get_config(self):
        config = super(VAE, self).get_config()
        # 序列化子模型
        config.update({
            "encoder": keras.saving.serialize_keras_object(self.encoder),
            "decoder": keras.saving.serialize_keras_object(self.decoder),
        })
        return config

    # 新增类方法，用于加载时反序列化子模型
    @classmethod
    def from_config(cls, config):
        encoder = keras.saving.deserialize_keras_object(config.pop("encoder"))
        decoder = keras.saving.deserialize_keras_object(config.pop("decoder"))
        return cls(encoder, decoder, **config)

2. 保存与加载模型 #

修改后的模型可正常保存，加载时需先注册自定义类（Sampling和VAE）：

# 保存模型
vae.save("vae_model")

# 加载模型（新环境/脚本中）
# 先定义Sampling类和上述VAE类，再执行加载
loaded_vae = keras.models.load_model("vae_model", custom_objects={"Sampling": Sampling, "VAE": VAE})

# 验证加载后的模型可用性
test_output = loaded_vae(synthetic_array[:10])
print(test_output[0].shape, test_output[1].shape)

关键说明 #

• call()方法是Keras模型的核心前向入口，必须实现才能让框架识别模型计算图，进而支持保存。
• get_config和from_config确保子模型（encoder、decoder）能正确序列化/反序列化，避免加载时出错。
• 训练时调用self(data)而非直接调用encoder，保证前向逻辑统一，让Keras能正确追踪梯度与模型结构。

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