Unable to save model using WeightNormalization in TF2.1 / load saved model in TF2.2rc4

[ghost]

System information

• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Linux Ubuntu 18.04
• TensorFlow version and how it was installed (source or binary): 2.1 / 2.2rc4, both installed from pip
• TensorFlow-Addons version and how it was installed (source or binary): 0.9.1, installed from pip
• Python version: 3.7.5
• Is GPU used? (yes/no): no

Describe the bug

Building a model with a WeightNormalization layer and saving it as a tensorflow saved model throws this error under tensorflow 2.1:

ValueError: Attempted to save a function b'__inference_model_layer_call_and_return_conditional_losses_692' which references a symbolic Tensor Tensor("model/weight_normalization_1/compute_weights/mul:0", shape=(3, 3, 3, 128), dtype=float32) that is not a simple constant. This is not supported.

Saving works under tensorflow 2.2rc4, but subsequently trying to load the saved model then throws this error:

KeyError: '__inference_model_layer_call_and_return_conditional_losses_516'

Code to reproduce the issue

import tensorflow as tf
import tensorflow_addons as tfa

# some basic model
inp = tf.keras.layers.Input((112,112,3))
x = tfa.layers.WeightNormalization(tf.keras.layers.Conv2D(128, kernel_size=3))(inp)
x = tf.keras.layers.GlobalAveragePooling2D()(x)
x = tf.keras.layers.Dense(2)(x)

model = tf.keras.models.Model(inputs=[inp], outputs=[x])

tf.saved_model.save(model, "./test_export") # throws in TF2.1
restored_model = tf.saved_model.load("./test_export") # throws in TF2.2rc4

[ghost]

Maybe this helps narrow it down further: I also just tried loading the model in TF2.2rc4 using

restored_model = tf.keras.models.load_model("./test_export")

the error message is now:

tensorflow.python.saved_model.nested_structure_coder.NotEncodableError: No encoder for object [] of type [<class 'tensorflow.core.protobuf.struct_pb2.StructuredValue'>].

[bhack]

Do you have the same error with model.save('model',save_format='tf')?

[ghost]

Yes

[ghost]

I can save/restore it using the hdf5 format, i.e. model.save(model, save_format="h5")but I need it in the tensorflow format for tensorflow serving.

[bhack]

Is It a subclassed model?

[ghost]

Yes in my case it is, but the sample code above also fails using just the vanilla keras model class

[bhack]

It could be useful if you can open a small PR creating a test like this one that it can cover your case: https://github.com/tensorflow/addons/blob/master/tensorflow_addons/layers/tests/wrappers_test.py#L149

Lets the CI on the PR to fail you test case

[ghost]

Cool, done

[bhack]

Thanks

[schmidt-jake]

I'm also having this issue on tensorflow 2.3.