Implementing custom convolutional layer in Keras - error when loading model

Question

I have implemented a minimal example of Wavenet, closely following the steps from here - https://github.com/basveeling/wavenet.

The issue is, that the model uses a custom layer, which works fine during training but once the model is reloaded, Keras cannot find the Causal Layer, even though I am using custom objects.

I am using tensorflow 1.13 and keras 2.2.4

Here is an example of the first three key/value pairs for objects.

objects = {'initial_causal_conv': <class 'wavenet_utils.CausalConv1D'>,
           'dilated_conv_1_tanh_s0': <class 'wavenet_utils.CausalConv1D'>,
           'dilated_conv_1_sigm_s0': <class 'wavenet_utils.CausalConv1D'>,
           '...': <class 'wavenet_utils.CausalConv1D'>,
           '...': <class 'wavenet_utils.CausalConv1D'>}

model.fit(x=[x_tr1, x_tr2],
             y=y_tr1,
             epochs=epochs,
             batch_size=batch_size,
             validation_data=([x_vl1, x_vl2], y_vl1),
             callbacks=[checkpoint, early_stopping],
             verbose=verbose,
             shuffle=True,
             class_weight=class_weight)
model = load_model('model.h5', custom_objects=objects)

Which then returns this error:

Traceback (most recent call last):
  File "/home/xxx/PycharmProjects/WAVE/DATA_NN.py", line 48, in <module>
    objects=objects)
  File "/home/xxx/PycharmProjects/WAVE/functions.py", line 572, in run_neural_net
    model = load_model('model_conv.h5', custom_objects=objects)
  File "/home/xxx/PycharmProjects/WAVE/venv/lib/python3.6/site-packages/keras/engine/saving.py", line 419, in load_model
    model = _deserialize_model(f, custom_objects, compile)
  File "/home/xxx/PycharmProjects/WAVE/venv/lib/python3.6/site-packages/keras/engine/saving.py", line 225, in _deserialize_model
    model = model_from_config(model_config, custom_objects=custom_objects)
  File "/home/xxx/PycharmProjects/WAVE/venv/lib/python3.6/site-packages/keras/engine/saving.py", line 458, in model_from_config
    return deserialize(config, custom_objects=custom_objects)
  File "/home/xxx/PycharmProjects/WAVE/venv/lib/python3.6/site-packages/keras/layers/__init__.py", line 55, in deserialize
    printable_module_name='layer')
  File "/home/xxx/PycharmProjects/WAVE/venv/lib/python3.6/site-packages/keras/utils/generic_utils.py", line 145, in deserialize_keras_object
    list(custom_objects.items())))
  File "/home/xxx/PycharmProjects/WAVE/venv/lib/python3.6/site-packages/keras/engine/network.py", line 1022, in from_config
    process_layer(layer_data)
  File "/home/xxx/PycharmProjects/WAVE/venv/lib/python3.6/site-packages/keras/engine/network.py", line 1008, in process_layer
    custom_objects=custom_objects)
  File "/home/xxx/PycharmProjects/WAVE/venv/lib/python3.6/site-packages/keras/layers/__init__.py", line 55, in deserialize
    printable_module_name='layer')
  File "/home/xxx/PycharmProjects/WAVE/venv/lib/python3.6/site-packages/keras/utils/generic_utils.py", line 138, in deserialize_keras_object
    ': ' + class_name)
ValueError: Unknown layer: CausalConv1D

When building the model, CausalConv1D must be imported from wavenet_utils.py

Below is the full build_model function And here is wavenet_utils, containing the class CausalConv1D:

from keras.layers import Conv1D
from keras.utils.conv_utils import conv_output_length
import tensorflow as tf


class CausalConv1D(Conv1D):
    def __init__(self, filters, kernel_size, init='glorot_uniform', activation=None,
                 padding='valid', strides=1, dilation_rate=1, bias_regularizer=None,
                 activity_regularizer=None, kernel_constraint=None, bias_constraint=None, use_bias=True, causal=False,
                 output_dim=1,
                 **kwargs):
        self.output_dim = output_dim

        super(CausalConv1D, self).__init__(filters,
                                           kernel_size=kernel_size,
                                           strides=strides,
                                           padding=padding,
                                           dilation_rate=dilation_rate,
                                           activation=activation,
                                           use_bias=use_bias,
                                           kernel_initializer=init,
                                           activity_regularizer=activity_regularizer,
                                           bias_regularizer=bias_regularizer,
                                           kernel_constraint=kernel_constraint,
                                           bias_constraint=bias_constraint,
                                           **kwargs)

        self.causal = causal
        if self.causal and padding != 'valid':
            raise ValueError("Causal mode dictates border_mode=valid.")

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

    def call(self, x):
        if self.causal:
            def asymmetric_temporal_padding(x, left_pad=1, right_pad=1):
                pattern = [[0, 0], [left_pad, right_pad], [0, 0]]
                return tf.pad(x, pattern)

            x = asymmetric_temporal_padding(x, self.dilation_rate[0] * (self.kernel_size[0] - 1), 0)
        return super(CausalConv1D, self).call(x)

    def compute_output_shape(self, input_shape):
        input_length = input_shape[1]

        if self.causal:
            input_length += self.dilation_rate[0] * (self.kernel_size[0] - 1)

        length = conv_output_length(input_length,
                                    self.kernel_size[0],
                                    self.padding,
                                    self.strides[0],
                                    dilation=self.dilation_rate[0])

        shape = tf.TensorShape(input_shape).as_list()
        shape[-1] = self.output_dim
        return (input_shape[0], length, self.filters)

    def get_config(self):
        base_config = super(CausalConv1D, self).get_config()
        base_config['output_dim'] = self.output_dim
        return base_config

EDIT:

I have tried this approach before as well.

objects = {'CausalConv1D': <class 'wavenet_utils.CausalConv1D'>}

model.fit(x=[x_tr1, x_tr2],
             y=y_tr1,
             epochs=epochs,
             batch_size=batch_size,
             validation_data=([x_vl1, x_vl2], y_vl1),
             callbacks=[checkpoint, early_stopping],
             verbose=verbose,
             shuffle=True,
             class_weight=class_weight)
model = load_model('model.h5', custom_objects=objects)

Which then returns this error:

Traceback (most recent call last):
  File "/home/xxx/PycharmProjects/WAVE/DATA_NN.py", line 47, in <module>
    objects=objects)
  File "/home/xxx/PycharmProjects/WAVE/functions.py", line 574, in run_neural_net
    model = load_model('model.h5', custom_objects=objects)
  File "/home/xxx/PycharmProjects/WAVE/venv/lib/python3.6/site-packages/keras/engine/saving.py", line 419, in load_model
    model = _deserialize_model(f, custom_objects, compile)
  File "/home/xxx/PycharmProjects/WAVE/venv/lib/python3.6/site-packages/keras/engine/saving.py", line 225, in _deserialize_model
    model = model_from_config(model_config, custom_objects=custom_objects)
  File "/home/xxx/PycharmProjects/WAVE/venv/lib/python3.6/site-packages/keras/engine/saving.py", line 458, in model_from_config
    return deserialize(config, custom_objects=custom_objects)
  File "/home/xxx/PycharmProjects/WAVE/venv/lib/python3.6/site-packages/keras/layers/__init__.py", line 55, in deserialize
    printable_module_name='layer')
  File "/home/xxx/PycharmProjects/WAVE/venv/lib/python3.6/site-packages/keras/utils/generic_utils.py", line 145, in deserialize_keras_object
    list(custom_objects.items())))
  File "/home/xxx/PycharmProjects/WAVE/venv/lib/python3.6/site-packages/keras/engine/network.py", line 1022, in from_config
    process_layer(layer_data)
  File "/home/xxx/PycharmProjects/WAVE/venv/lib/python3.6/site-packages/keras/engine/network.py", line 1008, in process_layer
    custom_objects=custom_objects)
  File "/home/xxx/PycharmProjects/WAVE/venv/lib/python3.6/site-packages/keras/layers/__init__.py", line 55, in deserialize
    printable_module_name='layer')
  File "/home/xxx/PycharmProjects/WAVE/venv/lib/python3.6/site-packages/keras/utils/generic_utils.py", line 147, in deserialize_keras_object
    return cls.from_config(config['config'])
  File "/home/xxx/PycharmProjects/WAVE/venv/lib/python3.6/site-packages/keras/engine/base_layer.py", line 1109, in from_config
    return cls(**config)
  File "/home/xxx/PycharmProjects/WAVE/wavenet_utils.py", line 26, in __init__
    **kwargs)
TypeError: __init__() got multiple values for keyword argument 'kernel_initializer'

Could this be the issue mentioned here https://github.com/keras-team/keras/issues/12316?

And if so, is there any way around it?

Answer 1

There is only one custom object, which is CausalConv1D.

objects = {'CausalConv1D': wavenet_utils.CausalConv1D}

Now you must be sure that your get_config method is correct and has everything needed in the __init__ method of your layer.

It misses the causal property and has a kernel_initializer coming from the base class that is not supported by your __init__ method.

Let's list every property you need, and then check which ones are in the base config:

• filters: in base
• kernel_size: in base
• init: not in base, but there is kernel_initializer in base!!!!!
  • kernel_initializer is a config item that your __init__ method doesn't support
  • rename this init parameter to kernel_initializer
• activation: in base
• padding: in base
• strides: in base
• dilation_rate: in base
• bias_regularizer: in base
• activity_regularizer: in base
• kernel_constraint: in base
• bias_constraint: in base
• use_bias: in base
• causal: not in base!
  • must add this in your config! (or the model will always use the default value)
• output_dim: not in base!
• **kwargs: in base

Layer's __init__:

def __init__(self, filters, kernel_size, 

             ############## here:
             kernel_initializer='glorot_uniform', 
             #############

             activation=None,
             padding='valid', strides=1, dilation_rate=1, bias_regularizer=None,
             activity_regularizer=None, kernel_constraint=None, bias_constraint=None, use_bias=True, causal=False,
             output_dim=1,
             **kwargs):

Layer's get_config

It must contain all __init__ params that are not in the base class:

def get_config(self):
    base_config = super(CausalConv1D, self).get_config()
    base_config['causal'] = self.causal
    base_config['output_dim'] = self.output_dim
    return base_config

Answer 2

Somehow, no approach I've tried so far has been able to correctly load the model when using load_model. Below is a simple work around which only saves the weights, then deletes the existing model, builds a new one and compiles it again, and loads saved the weights which do save correctly, even with custom layers present.

model = build_model()

checkpoint = ModelCheckpoint('model.h5', monitor='val_acc',
                             verbose=1, save_best_only=True, save_weights_only=True, mode='max')

model.fit(x, y)

del model

model = build_model()

model.load_weights('model.h5')

model.predict(x_test)