使用自定义图层保存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是使用Custom_Objects，而loading Saved Model，即替换代码，

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之前，因此在Loading时应使用Custom Objects.

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

更正数字2是在自定义层的__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屋！