求助：带自定义层的Keras预训练模型无法保存为.h5格式

核心原因 #

HDF5（.h5）格式保存模型时，要求所有自定义层都能被序列化，而自定义层默认未实现get_config()方法，导致无法正确保存层的配置信息，进而触发保存失败。

解决方案 #

方案1：给自定义层添加get_config()方法 #

在自定义层类中必须实现get_config()方法，返回层的初始化参数，让Keras能正确序列化该层。示例代码：

class CustomLayer(tf.keras.layers.Layer):
    def __init__(self, units=32, **kwargs):
        super(CustomLayer, self).__init__(**kwargs)
        self.units = units

    def build(self, input_shape):
        self.kernel = self.add_weight(shape=(input_shape[1], self.units),
                                      initializer='uniform',
                                      name='kernel')

    def call(self, inputs):
        return tf.matmul(inputs, self.kernel)

    # 必须添加该方法
    def get_config(self):
        config = super(CustomLayer, self).get_config()
        config.update({'units': self.units})
        return config

完成修改后，即可用model.save('model.h5')保存模型。加载时若自定义层不在Keras默认注册表，需指定custom_objects参数：

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

方案2：改用SavedModel格式保存 #

若不想修改自定义层代码，可直接使用TensorFlow推荐的SavedModel格式，它对自定义层的兼容性更好，无需额外处理序列化：

model.save('my_saved_model')  # 保存为文件夹形式

加载时直接调用：

loaded_model = tf.keras.models.load_model('my_saved_model')

方案3：只保存模型权重，再重新构建模型加载 #

如果只需要保留模型参数，可以先保存权重文件，之后重新构建与训练时一致的模型结构，再加载权重：

# 保存权重
model.save_weights('model_weights.h5')

# 加载流程：先重新构建模型结构
base_model = tf.keras.applications.ResNet50(weights='imagenet', include_top=False)
# 重新添加你的自定义层...
model = tf.keras.Model(inputs=base_model.input, outputs=custom_output)
# 加载权重
model.load_weights('model_weights.h5')

额外注意事项 #

• 若自定义层包含复杂的非张量操作，可能还需要实现from_config()方法，但多数场景下get_config()即可满足需求。
• 优先选择SavedModel格式，它对自定义层、自定义训练循环的支持更完善，是TensorFlow官方推荐的模型保存方式。

内容的提问来源于stack exchange，提问作者Youvraj Singh Gaur