Serialize custom dynamic layer in keras tensorflow

Question

I am rather new to tensorflow as well as Python (transfering from R). Currently I am working on a recommendation system in python using keras and tensorflow. The data is "unary", so I only know if someone clicked on something or if he didn't.

The core model is build with the functional API and is a basic wide model (input = multi-hot encoded binary rating matrix, output = probability per class) with one hidden layer in between. To make the model easier to use, I want to be able to throw in class labels and output topN predictions with class labels and probabilities. So instead of "Input = [[0,1,0,1,0,0,0,0]], Ouput = [[0.3,0.1,0.6,0.0]" I want to be able to do like "Input = [['apple','orange','bean']], Ouput = [['lemon',banana'],[0.3,0.2]]".

To do this, I train the basic model and then wrap two custom layers around the model, one at the beginning and one at the end (like here: https://towardsdatascience.com/customize-classification-model-output-layer-46355a905b86 ) (I also tried feature_columns, they did not really do it for me). To make the input custom layer, I had to set the layer to "dynamic = True", to enable eager execution. I could not find a way to create a layer for this "tokenization" without using eager execution. This works fine so far.

But now I can not restore the saved model (neither using h5 nor save_model.pb). I also specified the "get_config" method for the custom layers, and everything works fine as long as I only safe the model with the second custom layer at the end. So I suppose the error occurs because the first layer is dynamic. So how do I serealize a dynamic custom layer in keras?

I would really appreciate any help or even thoughts as I could not find any matching topic (or even any topic covering dynamic custom layers in general). Please find some code here:

class LabelInputLayer(Layer):

    def __init__(self, labels, n_labels, **kwargs):
        self.labels = labels
        self.n_labels = n_labels
        super(LabelInputLayer, self).__init__(**kwargs)

    def call(self, x):
        batch_size = tf.shape(x)[0]

        tf_labels = tf.constant([self.labels], dtype="int32")
        n_labels = self.n_labels
        x = tf.constant(x)
        tf_labels = tf.tile(tf_labels,[batch_size,1])

# go through every instance and multi-hot encode labels

        for j in tf.range(batch_size):
            index = []
            for i in tf.range(n_labels):
                if tf_labels[j,i].numpy() in x[j,:].numpy():
                    index.append(True)
                else:
                    index.append(False)

# create initial rating matrix and append for each instance

            if j == 0:
                rating_te = tf.where(index,1,0)
            else:
                rating_row = tf.where(index,1,0)
                rating_te = tf.concat([rating_te,rating_row],0)

        rating_te = tf.reshape(rating_te, [batch_size,-1])

        return [rating_te]

    def compute_output_shape(self, input_shape):
        return tf.TensorShape([None, self.n_labels])

# define get_config to enable serialization of the layer

    def get_config(self):
        config={'labels':self.labels,
                'n_labels':self.n_labels}
        base_config = super(LabelInputLayer, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

Here I am creating the basic model:

#Create Functional Model---------------------

a = tf.keras.layers.Input(shape=[n_classes])
b = tf.keras.layers.Dense(4000, activation = "relu", input_dim = n_classes)(a)
c = tf.keras.layers.Dropout(rate = 0.2)(b)
d = tf.keras.layers.Dense(n_classes, activation = "softmax")(c)
nn_model = tf.keras.Model(inputs = a, outputs = d)

And this is the final model (works in the notebook but can not be restored once saved, so no use in production):

input_layer = Input(shape = input_shape, name = "first_input")
encode_layer = LabelInputLayer(labels = labels, n_labels = n_labels, dynamic = True, input_shape = input_shape)(input_layer)
pre_trained = tf.keras.models.Sequential(nn_model.layers[1:])(encode_layer)
decode_layer = LabelLimitLayer(labels, n_preds)(pre_trained)
encoder_model = tf.keras.Model(inputs = input_layer, outputs = decode_layer)

Save and restore the model:

tf.saved_model.save(encoder_model, "encoder_model")

model = tf.keras.models.load_model("encoder_model")

This is the error I receive if I want to restore the model in another notebook (Unfortunately I can also not use the "custom_objects" parameter in the load method, as I have to deploy the model from the save file only):

ValueError: Could not find matching function to call loaded from the SavedModel. Got:
  Positional arguments (1 total):
    * Tensor("x:0", shape=(None, 10), dtype=float32)
  Keyword arguments: {}

Expected these arguments to match one of the following 0 option(s):

Answer 1

I managed to get around the Issue by using tf.functions instead of dynamic=True. I did not however manage to save a dynamic layer in Tensorflow. Maybe that helps someone.