如何从 .h5 文件正确加载带有自定义层的 Keras 模型？

Question

I built a Keras model with a custom layers, and it was saved to a .h5 file by the callback ModelCheckPoint . When I tried to load this model after the training, the error message below showed up:

 __init__() missing 1 required positional argument: 'pool_size'

This is the definition of the custom layer and its __init__ method:

class MyMeanPooling(Layer):
    def __init__(self, pool_size, axis=1, **kwargs):
        self.supports_masking = True
        self.pool_size = pool_size
        self.axis = axis
        self.y_shape = None
        self.y_mask = None
        super(MyMeanPooling, self).__init__(**kwargs)

This is how I add this layer to my model:

x = MyMeanPooling(globalvars.pool_size)(x)

This is how I load the model:

from keras.models import load_model

model = load_model(model_path, custom_objects={'MyMeanPooling': MyMeanPooling})

These are the full error messages:

Traceback (most recent call last):
  File "D:/My Projects/Attention_BLSTM/script3.py", line 9, in <module>
    model = load_model(model_path, custom_objects={'MyMeanPooling': MyMeanPooling})
  File "D:\ProgramData\Anaconda3\envs\tf\lib\site-packages\keras\engine\saving.py", line 419, in load_model
    model = _deserialize_model(f, custom_objects, compile)
  File "D:\ProgramData\Anaconda3\envs\tf\lib\site-packages\keras\engine\saving.py", line 225, in _deserialize_model
    model = model_from_config(model_config, custom_objects=custom_objects)
  File "D:\ProgramData\Anaconda3\envs\tf\lib\site-packages\keras\engine\saving.py", line 458, in model_from_config
    return deserialize(config, custom_objects=custom_objects)
  File "D:\ProgramData\Anaconda3\envs\tf\lib\site-packages\keras\layers\__init__.py", line 55, in deserialize
    printable_module_name='layer')
  File "D:\ProgramData\Anaconda3\envs\tf\lib\site-packages\keras\utils\generic_utils.py", line 145, in deserialize_keras_object
    list(custom_objects.items())))
  File "D:\ProgramData\Anaconda3\envs\tf\lib\site-packages\keras\engine\network.py", line 1022, in from_config
    process_layer(layer_data)
  File "D:\ProgramData\Anaconda3\envs\tf\lib\site-packages\keras\engine\network.py", line 1008, in process_layer
    custom_objects=custom_objects)
  File "D:\ProgramData\Anaconda3\envs\tf\lib\site-packages\keras\layers\__init__.py", line 55, in deserialize
    printable_module_name='layer')
  File "D:\ProgramData\Anaconda3\envs\tf\lib\site-packages\keras\utils\generic_utils.py", line 147, in deserialize_keras_object
    return cls.from_config(config['config'])
  File "D:\ProgramData\Anaconda3\envs\tf\lib\site-packages\keras\engine\base_layer.py", line 1109, in from_config
    return cls(**config)
TypeError: __init__() missing 1 required positional argument: 'pool_size'

Answer 1

Actually I don't think you can load this model.

The most likely issue is that you did not implement the get_config() method in your layer. This method returns a dictionary of configuration values that should be saved:

def get_config(self):
    config = {'pool_size': self.pool_size,
              'axis': self.axis}
    base_config = super(MyMeanPooling, self).get_config()
    return dict(list(base_config.items()) + list(config.items()))

You have to retrain the model after adding this method to your layer, as the previously saved model does not have the configuration for this layer saved into it. This is why you cannot load it, it requires retraining after making this change.

Answer 2

From the answer of "LiamHe commented on Sep 27, 2017" on the following issue: https://github.com/keras-team/keras/issues/4871 .

I met the same problem today : ** TypeError: init() missing 1 required positional arguments**. Here is how I solve the problem : (Keras 2.0.2)

1. Give the positional arguments of the layer with some default values
2. Override get_config function to the layer with some thing like

def get_config(self):
    config = super().get_config()
    config['pool_size'] = # say self._pool_size  if you store the argument in __init__
    return config

3. Add layer class to custom_objects when you are loading model.

Answer 3

If you don't have enough time to retrain the model in the solution way of Matias Valdenegro. You can set the default value of pool_size in class MyMeanPooling like the following code. Note that the value of pool_size should be consistent with the value while training the model. Then you can load the model.

class MyMeanPooling(Layer):
    def __init__(self, pool_size, axis=1, **kwargs):
        self.supports_masking = True
        self.pool_size = 2  # The value should be consistent with the value while training the model
        self.axis = axis
        self.y_shape = None
        self.y_mask = None
        super(MyMeanPooling, self).__init__(**kwargs)

ref: https://www.jianshu.com/p/e97112c34e43