TensorFlow中Keras v3格式自定义VAE模型保存加载报错求助

问题根源 #

保存自定义Keras模型时，TensorFlow仅保存模型的结构参数与权重，不会自动保存自定义类的定义逻辑。加载时若未明确告知TensorFlow自定义类的位置，即便实现了get_config()方法（仅保证参数序列化），仍会触发找不到类的TypeError。

可行解决方案 #

1. 加载时显式指定自定义类 #

调用load_model时，通过custom_objects参数传入VariationalAutoEncoder类，让TensorFlow定位到类定义：

from tensorflow.keras.models import load_model

# 确保VariationalAutoEncoder类的代码在当前运行环境中可见（已定义或导入）
loaded_vae = load_model("vae211.keras", custom_objects={"VariationalAutoEncoder": VariationalAutoEncoder})

2. 给自定义类添加序列化注册装饰器 #

在VariationalAutoEncoder类定义上方添加@keras.saving.register_keras_serializable装饰器，让TensorFlow保存时自动记录类的注册信息，加载时无需额外参数：

import tensorflow as tf
from tensorflow import keras

@keras.saving.register_keras_serializable()
class VariationalAutoEncoder(keras.Model):
    def __init__(self, latent_dim=2, **kwargs):
        super().__init__(**kwargs)
        self.latent_dim = latent_dim
        # 编码器、解码器的定义代码...
    
    def get_config(self):
        config = super().get_config()
        config["latent_dim"] = self.latent_dim
        return config
    
    # 模型的call、sample等方法实现...

注意：此方法需修改类后重新执行vae.save()保存模型，之前的旧模型仍需用方案1加载。

3. 确保加载时类定义在作用域内 #

若自定义类写在单独模块文件（如vae_models.py），加载模型前先导入该类：

from vae_models import VariationalAutoEncoder
loaded_vae = load_model("vae211.keras")

关键提示 #

• get_config()是必要的，它保证自定义类的初始化参数能正确序列化/反序列化，但无法单独解决类的定位问题。
• 若模型包含自定义层，需对自定义层做同样处理（加注册装饰器或在custom_objects中指定）。

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