15

I have defined a custom layer named CapsuleLayer. The actual model has been defined in a separate class. I have loaded the weights into the actual model and have saved the model in an .h5 file. However when I try to load the model using load_model(filepath) I get the error

ValueError: Unknown layer: CapsuleLayer

How can I incorporate the custom layer into my model while loading the saved model.

benjaminplanche
  • 14,689
  • 5
  • 57
  • 69
Sounak Ray
  • 303
  • 1
  • 3
  • 9

2 Answers2

28

C.f. Keras FAQ, "Handling custom layers (or other custom objects) in saved models":

If the model you want to load includes custom layers or other custom classes or functions, you can pass them to the loading mechanism via the custom_objects argument:

from keras.models import load_model
# Assuming your model includes instance of an "AttentionLayer" class
model = load_model('my_model.h5', custom_objects={'AttentionLayer': AttentionLayer})

Alternatively, you can use a custom object scope:

from keras.utils import CustomObjectScope

with CustomObjectScope({'AttentionLayer': AttentionLayer}):
    model = load_model('my_model.h5')

Custom objects handling works the same way for load_model, model_from_json, model_from_yaml:

from keras.models import model_from_json
model = model_from_json(json_string, custom_objects={'AttentionLayer': AttentionLayer})

In your case, model = load_model('my_model.h5', custom_objects={'CapsuleLayer': CapsuleLayer}) should solve your problem.

benjaminplanche
  • 14,689
  • 5
  • 57
  • 69
  • I need to import the class as well I guess. – Sounak Ray Jun 13 '18 at 17:16
  • Exact. The `custom_objects` parameter basically has for purpose to pass your custom class and its name to `load_model()`, so that it can recognize and recover it from the saved files. – benjaminplanche Jun 13 '18 at 17:19
  • I am getting an error. TypeError: __init__() takes atleast 3 arguments (1 given). What could be the problem.? – Sounak Ray Jun 13 '18 at 17:34
  • 1
    Your custom layer probably needs 3 arguments to be initialized, but doesn't receive them. You may have to add a `get_config()` method to it. This [issue's answer](https://github.com/keras-team/keras/issues/1927#issuecomment-309712640) may help you, as well as the code for the core layers e.g. [here](https://github.com/keras-team/keras/blob/master/keras/layers/convolutional.py#L214). – benjaminplanche Jun 14 '18 at 13:23
6

Just for completeness I am adding just a bit on top of the answer of benjaminplanche. If your custom layer AttentionLayer have any initial parameter that configure its behaviour you need to implement the get_config method of the class. Otherwise it will fail to load. I am writing this because I had a lot of troubles on how to load custom layers with arguments, so I'll leave it here.

For example, a dummy implementation of your layer:

class AttentionLayer(Layer):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def build(self, input_shape):
        return super().build(input_shape)

    def call(self, x):
        # Implementation about how to look with attention!
        return x

    def compute_output_shape(self, input_shape):
        return input_shape

This will load with any of the approaches detailed in benjaminplanche's answer, i.e. using the custom_objects={'AttentionLayer': AttentionLayer}. However if your layer have some arguments the loading would fail.

Imagine the init method of your class have 2 paramters:

class AttentionLayer(Layer):
    def __init__(self, param1, param2, **kwargs):
        self.param1 = param1
        self.param2 = param2
        super().__init__(**kwargs)

Then, when you load it with:

model = load_model('my_model.h5', custom_objects={'AttentionLayer': AttentionLayer})

It would throw this error:

Traceback (most recent call last):
  File "/path/to/file/cstm_layer.py", line 62, in <module>
    h = AttentionLayer()(x)
TypeError: __init__() missing 2 required positional arguments: 'param1' and 'param2'

In order to solve it you need to implement the get_config method in your custom layer class. An example:

class AttentionLayer(Layer):
    def __init__(self, param1, param2, **kwargs):
        self.param1 = param1
        self.param2 = param2
        super().__init__(**kwargs)

    # ...

    def get_config(self):
        # For serialization with 'custom_objects'
        config = super().get_config()
        config['param1'] = self.param1
        config['param2'] = self.param2
        return config

So when you save the model, the saving routine will call the get_config and will serialize the inner state of your custom layer, i.e., the self.params. And when you load it, the loader will know how to initialize the inner state of your custom layer.

JVGD
  • 657
  • 9
  • 15
  • what are idx_init and idx_end please – nassim May 15 '20 at 16:42
  • Fixed, it is param1 and 2. In the get_config method you should store any Layer attributes relevant for its operation (in order to load the layer later with the correct conf) – JVGD May 16 '20 at 11:17