I have saved the model by disabling eager execution in tf2.6 version. Below is the snippet used for saving and running

import tensorflow as tf
tf.compat.v1.disable_eager_execution()

x = tf.keras.Input(shape=(224, 224, 3), batch_size=None)
x1=tf.keras.Input(1, dtype=tf.int32)
print(type(x1), type(x)) #<class 'tensorflow.python.framework.ops.Tensor'> <class 'tensorflow.python.framework.ops.Tensor'>
y = tf.cond(tf.less(x1,5), lambda :tf.keras.layers.ReLU()(x), lambda :tf.keras.layers.LeakyReLU(alpha=0.1)(x))
model=tf.keras.models.Model(inputs=[x,x1], outputs=[y])
model.save("cond.h5")
model.summary()

Running

import tensorflow as tf
import numpy as np
from tensorflow.keras.models import load_model
tf.compat.v1.disable_eager_execution()


model=load_model('cond.h5')
x = np.ones((1, 224,224,3), dtype=np.float32)
x1 = np.zeros((1, 1), dtype=np.int32)
z=model.predict([x,x1], batch_size=1)

Error:

tensorflow.python.framework.errors_impl.InvalidArgumentError: Op type not registered 'cond_true_5' in binary running on anilm. Make sure the Op and Kernel are registered in the binary running in this process. Note that if you are loading a saved graph which used ops from tf.contrib, accessing (e.g.) tf.contrib.resampler should be done before importing the graph, as contrib ops are lazily registered when the module is first accessed. while building NodeDef 'tf_op_layer_cond/cond/then/_0'

CodePudding user response：

I think the tf.cond operation is causing a problem because it is not a Keras layer. You could try wrapping the operation in a Lambda layer:

import tensorflow as tf

tf.compat.v1.disable_eager_execution()

x = tf.keras.Input(shape=(224, 224, 3), batch_size=None)
x1 = tf.keras.Input(1, dtype=tf.int32)
y = tf.keras.layers.Lambda(lambda x: tf.cond(tf.less(x[0], 5), lambda: tf.keras.layers.ReLU()(x[1]), lambda: tf.keras.layers.LeakyReLU(alpha=0.1)(x[1])))([x1, x])
model = tf.keras.models.Model(inputs=[x, x1], outputs=[y])
model.save("cond.h5")
model.summary()

I would still recommend a custom layer though.