使用自定义层保存 Keras 模型

Question

我正在尝试将 Keras 模型保存在 H5 文件中。 Keras 模型有一个自定义层。 当我尝试恢复模型时，我收到以下错误：

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-5-0fbff9b56a9d> in <module>()
      1 model.save('model.h5')
      2 del model
----> 3 model = tf.keras.models.load_model('model.h5')

8 frames
/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/utils/generic_utils.py in class_and_config_for_serialized_keras_object(config, module_objects, custom_objects, printable_module_name)
    319   cls = get_registered_object(class_name, custom_objects, module_objects)
    320   if cls is None:
--> 321     raise ValueError('Unknown ' + printable_module_name + ': ' + class_name)
    322 
    323   cls_config = config['config']

ValueError: Unknown layer: CustomLayer

能否请您告诉我应该如何保存和加载所有自定义 Keras 层的权重？ （另外，保存时没有提示，能不能从我已经保存但现在无法加载的H5文件中加载模型？）

这是针对此错误的最小工作代码示例 (MCVE)，以及完整的扩展消息：Google Colab Notebook

为了完整起见，这是我用来制作自定义层的代码。 get_config 和 from_config 都工作正常。

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.save('model.h5')
model = tf.keras.models.load_model('model.h5')

Answer 1

更正1是使用Custom_Objects而loadingSaved Model，即替换代码，

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

与

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

由于我们使用Custom Layers 到build 和Model 以及在Saving 它之前，我们应该使用Custom Objects 而Loading 它。

更正2是在Custom Layer的__init__函数中添加**kwargslike

def __init__(self, k, name=None, **kwargs):
        super(CustomLayer, self).__init__(name=name)
        self.k = k
        super(CustomLayer, self).__init__(**kwargs)

完整的工作代码如下所示：

import tensorflow as tf

class CustomLayer(tf.keras.layers.Layer):
    def __init__(self, k, name=None, **kwargs):
        super(CustomLayer, self).__init__(name=name)
        self.k = k
        super(CustomLayer, self).__init__(**kwargs)


    def get_config(self):
        config = super(CustomLayer, self).get_config()
        config.update({"k": self.k})
        return config

    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')
])
tf.keras.models.save_model(model, 'model.h5')
new_model = tf.keras.models.load_model('model.h5', custom_objects={'CustomLayer': CustomLayer})

print(new_model.summary())

以上代码的输出如下所示：

WARNING:tensorflow:No training configuration found in the save file, so the model was *not* compiled. Compile it manually.
Model: "sequential_1"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
custom_layer_1 (CustomLayer) (None, 10)                0         
_________________________________________________________________
output_layer (Dense)         (None, 1)                 11        
=================================================================
Total params: 11
Trainable params: 11
Non-trainable params: 0

希望这会有所帮助。快乐学习！

Answer 2

您可以在答案https://stackoverflow.com/a/62326857/8056572 中提到的load_model 方法中手动提供映射custom_objects，但是当您有很多自定义层（或定义的任何自定义可调用对象。例如指标、损失）时，这可能会很乏味，优化器，...）。

Tensorflow 提供了一个 utils 函数来自动完成：tf.keras.utils.register_keras_serializable

您必须按如下方式更新您的CustomLayer：

import tensorflow as tf

@tf.keras.utils.register_keras_serializable()
class CustomLayer(tf.keras.layers.Layer):
    def __init__(self, k, **kwargs):
        self.k = k
        super(CustomLayer, self).__init__(**kwargs)

    def get_config(self):
        config = super().get_config()
        config["k"] = self.k
        return config

    def call(self, input):
        return tf.multiply(input, 2)

---

这是完整的工作代码：

import tensorflow as tf


@tf.keras.utils.register_keras_serializable()
class CustomLayer(tf.keras.layers.Layer):
    def __init__(self, k, **kwargs):
        self.k = k
        super(CustomLayer, self).__init__(**kwargs)

    def get_config(self):
        config = super().get_config()
        config["k"] = self.k
        return config

    def call(self, input):
        return tf.multiply(input, 2)


def main():
    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')
        ]
    )
    print("SUMMARY OF THE MODEL CREATED")
    print("-" * 60)
    print(model.summary())
    model.save('model.h5')

    del model

    print()
    print()

    model = tf.keras.models.load_model('model.h5')
    print("SUMMARY OF THE MODEL LOADED")
    print("-" * 60)
    print(model.summary())

if __name__ == "__main__":
    main()

以及对应的输出：

SUMMARY OF THE MODEL CREATED
------------------------------------------------------------
Model: "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
custom_layer (CustomLayer)   (None, 10)                0         
_________________________________________________________________
output_layer (Dense)         (None, 1)                 11        
=================================================================
Total params: 11
Trainable params: 11
Non-trainable params: 0
_________________________________________________________________
None


WARNING:tensorflow:No training configuration found in the save file, so the model was *not* compiled. Compile it manually.
SUMMARY OF THE MODEL LOADED
------------------------------------------------------------
Model: "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
custom_layer (CustomLayer)   (None, 10)                0         
_________________________________________________________________
output_layer (Dense)         (None, 1)                 11        
=================================================================
Total params: 11
Trainable params: 11
Non-trainable params: 0
_________________________________________________________________
None