Why can't models with constants be saved and reloaded?

Keras
1

Does anyone know why this example can be serialized, but not deserialized?

import numpy as np
import tensorflow as tf

N=4
input = tf.keras.layers.Input(shape=[N])
constant=tf.constant(np.random.random([1,N]))
x = tf.keras.layers.Add()([input,constant])
m = tf.keras.Model(inputs=[input], outputs=[x])
print(m(np.random.random([1,N])))     # can run before saving
tf.keras.models.save_model(m,"test_dot.keras")
tf.keras.models.load_model("test_dot.keras")    # fail

The immediate error is in base_merge.py:
batch_sizes = {s[0] for s in input_shape if s} - {None}
TypeError: unhashable type: ‘list’

s is some unexpected value, [[(), (), (), ()]]

Another symptom is when you load the .keras into Netron and try to see the 2nd operand, it says “ERROR: Invalid tensor data length.” If the 2nd operand isn’t constant, then no issue. I’m using Tensorflow 2.13

3

Hi @Yale_Zhang

Issues

I’m not an expert but I think these are the issues:

  1. This input is a symbolic tensor, you can’t add an explicit add operation. You can add only another symbolic tensor.
  2. Normally you’d be able to achieve this with a single Lambda layer, like so:
def add_constant(x):
    add_value=np.random.random((1,N))
    return x + add_value
x = Lambda(add_constant)(input)

But from what I understand Lambda layers are discouraged (I think the reason is that they can’t be serializable, for which I understand saved.) and you will get an error.

Possible Solution

What I did was to use a custom layer, this is how I defined it (notice that I import directly from keras, I think this is the right way, but indeed tensorflow ships with a copy of it.).

The key step is to add:

@keras.saving.register_keras_serializable()
class MyAddLayer(keras.layers.Layer):
    def __init__(self):
        super().__init__()
    def call(self, x):
      return add_constant(x)

x = MyAddLayer()(input)

Full code (runs in notebook, I don’t share it because I don’t know whether it’s secure.)

import numpy as np
import tensorflow as tf
from keras.layers import Input, Add, Lambda
import keras

N=4
input = Input(shape=(N,))
print(input)

def add_constant(x):
    add_value=np.random.random((1,N))
    return x + add_value

@keras.saving.register_keras_serializable()
class MyAddLayer(keras.layers.Layer):
    def __init__(self):
        super().__init__()
    def call(self, x):
      return add_constant(x)

x = MyAddLayer()(input)

m = keras.models.Model(inputs=input, outputs=(x,))
print(m(np.random.random((1,N))))     # can run before saving
m.save("test_dot.keras")
keras.saving.load_model("test_dot.keras")    # fail

PS: I assumed you use keras 3, which as indicated in the docs should be used with

Replace from tensorflow import keras to import keras

5

Creating a custom layer seems to be the accepted solution, but I forgot to say that won’t work for me because the converter I use needs the elementary nodes and using custom layers hides it. I reached that conclusion because when I create the custom layer and save it, the call member is never called, so how can the connectivity be known.

It seems Keras’s disk format is not as expressive as the in memory format. I can save the above as a TF saved model. Also, ONNX can represent the above. It seems the most straight forward solution is to save it as TF, then convert to ONNX.

7

Official TensorFlow 2.16 + Python 3.12 – JARaaS Hybrid RAG - 6/17/2024
Note: Sources at the end of the response

The issue arises because TensorFlow’s model saving and loading mechanisms, particularly for Keras models, have some limitations when it comes to handling constants or variables that are not standard tensors within the computational graph.

Here’s what’s happening in your example:

  1. Model Definition:
    You have defined a Keras model with an input layer and added a constant tensor to it.

    input = tf.keras.layers.Input(shape=[N])
constant = tf.constant(np.random.random([1, N]))
x = tf.keras.layers.Add()([input, constant])
m = tf.keras.Model(inputs=[input], outputs=[x])

  2. Running the Model:
    The model runs successfully before saving:

    print(m(np.random.random([1, N])))

  3. Saving the Model:
    The model is saved without any issues:

    tf.keras.models.save_model(m, "test_dot.keras")

  4. Loading the Model:
    The failure occurs during model loading:

    tf.keras.models.load_model("test_dot.keras")    # fail

Explanation

The problem lies in how TensorFlow handles the constant when saving and loading the model. In TensorFlow, constants are static and can be serialized. However, during deserialization, the context and type of the constant (especially if it’s not a standard tensor but a more complex object like a list) become ambiguous.

Specifically, the error message:

TypeError: unhashable type: ‘list’

indicates that TensorFlow is encountering a list where it expects a different type. The list in question is likely coming from an intermediate representation of the constant during the loading process.

Possible Solution

To avoid this issue, you can redefine the constant tensor as a Keras layer, rather than directly using tf.constant. Here’s an alternative approach:

import numpy as np
import tensorflow as tf

N = 4
input = tf.keras.layers.Input(shape=[N])
constant_layer = tf.keras.layers.Lambda(lambda x: tf.constant(np.random.random([1, N])))(input)
x = tf.keras.layers.Add()([input, constant_layer])
m = tf.keras.Model(inputs=[input], outputs=[x])
print(m(np.random.random([1, N])))     # can run before saving
tf.keras.models.save_model(m, "test_dot.keras")
loaded_model = tf.keras.models.load_model("test_dot.keras")
print(loaded_model(np.random.random([1, N])))    # should run after loading now

Summary

By using a Lambda layer to wrap the constant, the constant becomes part of the model’s graph, which can be correctly serialized and deserialized by TensorFlow.

Sources

  • Save and Load Models: save_and_load.ipynb (internal document)
  • TensorFlow Model Serialization Guide: saved_model.ipynb (internal document)
9

@Yale_Zhang you may need to provide a complete example, i’m able to get quite complex keras 3 models to the browser using the converters.