addons

addons copied to clipboard

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