Tensorflow 1.14+ Serialize Subclassed Keras Layers?

Question

I've gone through and re-re-re-read the Tensorflow Keras docs, e.g.

I have a simple subclassed layer:

class SimpleLayer(tf.keras.layers.Layer):
    def __init__(self, filters, kernel_size, **kwargs):
        super(SimpleLayer, self).__init__()
        self.filters = filters
        self.kernel_size = kernel_size
        self.c1 = tf.keras.layers.Conv1D(filters, kernel_size, padding='same', activation='relu')
        self.c2 = tf.keras.layers.Conv1D(filters, kernel_size, padding='same')

    def call(self, inputs):
        x = inputs
        x = self.c1(x)
        x = self.c2(x)
        return x

    def get_config(self):
        # config = super(tf.keras.layers.Layer, self).get_config()
        config = {}
        config.update({
            'filters': self.filters,
            'kernel_size': self.kernel_size,
        })
        return config

and then have a functional model:


x = tf.keras.Inputs(...)

# some keras layers
y = tf.keras.layers... (x)

# my keras layer
y = SimpleLayer(...)(y)

# some keras layers
y = tf.keras.layers... (y)
y = tf.keras.layers.Dense(1)(y)

model = tf.keras.Model(inputs=x, outputs=y)
model.compile(...)

model.fit(...)

model.save('model.h5')

and figure then I could load the model as:

tf.keras.models.load_model('model.h5')

but I get:

ValueError: Unknown layer: SimpleLayer

from the docs:

If you need your custom layers to be serializable as part of a Functional model, you can optionally implement a get_config method

which I have.

What am I doing wrong?

Dr. Snoopy Dr. Snoopy · Accepted Answer · 2019-09-05T05:21:34

You need to tell keras about your custom layer during loading, you do this with the custom_objects parameter:

tf.keras.models.load_model('model.h5', custom_objects = {'SimpleLayer': SimpleLayer})

Tensorflow 1.14+ Serialize Subclassed Keras Layers?

1 Answers