含自定义层CustomLayer的Keras模型保存与加载失败问题求助

为什么会出现ValueError: Unknown layer: CustomLayer？ #

当你用model.save()保存H5模型时，文件里会存储模型的结构、权重和配置，但Keras序列化系统默认只认识内置层。对于自定义层，它只保存了层的配置字典（也就是你在get_config()里返回的内容），但加载时找不到对应的类定义，自然就会报错。

正确保存并加载带自定义层的模型 #

有两种实用方法可以解决这个问题：

方法1：加载时显式指定自定义层 #

这是最直接的方式，在调用load_model()时，把你的自定义层类传入custom_objects参数，告诉Keras这个层的定义在哪里：

model = tf.keras.models.load_model('model.h5', custom_objects={'CustomLayer': CustomLayer})

这样Keras就能通过名称匹配到对应的自定义层类，顺利还原模型的结构和权重。

方法2：注册自定义层（推荐长期使用） #

如果你不想每次加载都手动传custom_objects，可以给自定义层类加上@tf.keras.utils.register_keras_serializable()装饰器，把它注册到Keras的序列化系统中：

@tf.keras.utils.register_keras_serializable()
class CustomLayer(tf.keras.layers.Layer):
    def __init__(self, k, name=None):
        super(CustomLayer, self).__init__(name=name)
        self.k = k
    def get_config(self):
        return {'k': self.k}
    def call(self, input):
        return tf.multiply(input, 2)

之后保存和加载就可以像使用内置层一样直接操作，不需要额外参数：

model.save('model.h5')
model = tf.keras.models.load_model('model.h5')

如何修复已保存的无法打开的H5模型？ #

别担心，已经保存的H5文件完全可以修复！只要你还有自定义层的类定义（或者能写出和之前完全一致的类，包括__init__的参数、get_config()返回的键值对），就可以用方法1中的custom_objects参数加载。

举个例子，就算你之前保存时没注册层，现在只要写出和原来一样的CustomLayer类，运行下面的代码就能加载：

# 先定义好和之前完全一致的CustomLayer类
class CustomLayer(tf.keras.layers.Layer):
    def __init__(self, k, name=None):
        super(CustomLayer, self).__init__(name=name)
        self.k = k
    def get_config(self):
        return {'k': self.k}
    def call(self, input):
        return tf.multiply(input, 2)

# 加载时传入custom_objects
model = tf.keras.models.load_model('model.h5', custom_objects={'CustomLayer': CustomLayer})

加载成功后，你可以重新保存一次（比如用注册后的类），之后再加载就不需要传额外参数了。

完整可验证代码 #

下面是修改后的完整代码，确保保存和加载流程完全正常：

import tensorflow as tf

@tf.keras.utils.register_keras_serializable()
class CustomLayer(tf.keras.layers.Layer):
    def __init__(self, k, name=None):
        super(CustomLayer, self).__init__(name=name)
        self.k = k
    def get_config(self):
        return {'k': self.k}
    def call(self, input):
        return tf.multiply(input, 2)

# 构建模型
model = tf.keras.models.Sequential([
    tf.keras.Input(name='input_layer', shape=(10,)),
    CustomLayer(10, name='custom_layer'),
    tf.keras.layers.Dense(1, activation='sigmoid', name='output_layer')
])

# 编译（如果需要训练的话）
model.compile(optimizer='adam', loss='binary_crossentropy')

# 保存模型
model.save('model.h5')

# 删除当前模型实例
del model

# 加载模型
model = tf.keras.models.load_model('model.h5')

# 验证模型结构和自定义层参数
print(model.summary())
print("自定义层的k参数值：", model.get_layer('custom_layer').k)

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