Loading custom CTC layer from h5 file in Keras

Question

I have a CTCLayer class like this:

class CTCLayer(layers.Layer):
def __init__(self, name=None):
    super().__init__(name=name)
    self.loss_fn = keras.backend.ctc_batch_cost


def call(self, y_true, y_pred):
    # Compute the training-time loss value and add it
    # to the layer using `self.add_loss()`.
    batch_len = tf.cast(tf.shape(y_true)[0], dtype="int64")
    input_length = tf.cast(tf.shape(y_pred)[1], dtype="int64")
    label_length = tf.cast(tf.shape(y_true)[1], dtype="int64")

    input_length = input_length * tf.ones(shape=(batch_len, 1), dtype="int64")
    label_length = label_length * tf.ones(shape=(batch_len, 1), dtype="int64")

    loss = self.loss_fn(y_true, y_pred, input_length, label_length)
    self.add_loss(loss)

    # At test time, just return the computed predictions
    return y_pred

I trained my model, saved it to a model.h5 file and loaded it through:

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

It's throwing a init() got an unexpected keyword argument 'trainable' error.

Since I don't want to train my model again (time constraint), is there any workaround I can do to load the model without having to add a get_config() in the CTCLayer class?

And if not, how should I modify a get_config() in the class?

Answer 1

This should work:

    class CTCLayer(layers.Layer):
    def __init__(self, name=None):
        def __init__(self, name=None, **kwargs):
        super(CTCLayer, self).__init__(name=name, **kwargs)