使用自定义层保存 Keras 模型 [英] Saving Keras models with Custom Layers

问题描述

我正在尝试将 Keras 模型保存在 H5 文件中.Keras 模型有一个自定义层.当我尝试恢复模型时，出现以下错误:

I am trying to save a Keras model in a H5 file. The Keras model has a custom layer. When I try to restore the model, I get the following error:

---------------------------------------------------------------------------
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文件中加载模型?)

Could you please tell me how I am supposed to save and load weights of all the custom Keras layers too? (Also, there was no warning when saving, will it be possible to load models from H5 files which I have already saved but can't load back now?)

以下是此错误的最小工作代码示例 (MCVE)，以及完整的扩展消息:Google Colab 笔记本

Here is the minimal working code sample (MCVE) for this error, as well as the full expanded message: Google Colab Notebook

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

Just for completeness, this is the code I used to make my custom layer. get_config and from_config are both working fine.

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')

推荐答案

更正编号 1 是在 loading Saved Model 时使用 Custom_Objects即，替换代码，

Correction number 1 is to use Custom_Objects while loading the Saved Model i.e., replace the code,

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 它之前，我们应该使用 <代码>自定义对象而加载它.

Since we are using Custom Layers to build the Model and before Saving it, we should use Custom Objects while Loading it.

更正二是在Custom Layer的__init__函数中添加**kwargs，如

Correction number 2 is to add **kwargs in the __init__ function of the Custom Layer like

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

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

Complete working code is shown below:

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())

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

Output of the above code is shown below:

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

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

Hope this helps. Happy Learning!

这篇关于使用自定义层保存 Keras 模型的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！