ValueError: Unknown layer: CapsuleLayer

Question

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.

Answer 1

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.

Answer 2

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.