I'm fine tuning an EfficientNetB0 model (from Keras Applications in TF 2.9.2) and I want to optimize the final model as much as possible in order to be inferred via Tensorflow C++ API. Applying tf.saved_model.save after tensorflow.python.framework.convert_to_constants.convert_variables_to_constants_v2 works fine for several architectures such as MobileNet or ResNet50, but fails for EfficientNetBX family.
As a minimal code that illustrates the described problem, the following yields the expected results:
import tensorflow as tf
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
input_shape = (224, 224, 3)
model = tf.keras.applications.ResNet50(weights='imagenet', include_top=True, input_shape=input_shape)
#model = tf.keras.applications.EfficientNetB0(weights='imagenet', include_top=True, input_shape=input_shape)
full_model = tf.function(lambda x: model(x))
full_model = full_model.get_concrete_function(tf.TensorSpec(model.inputs[0].shape,
model.inputs[0].dtype,
name='yourInputName'))
frozen_func = convert_variables_to_constants_v2(full_model)
tf.saved_model.save(frozen_func, 'saved_model')
...but fails when switching to EfficientNetB0. More precisely, it throws the following error:
ValueError: Found invalid capture
Tensor("efficientnetb0/stem_activation/beta:0", shape=(), dtype=float32) when saving custom gradients
I would like to know if I am missing any particular step that is required for EfficientNetB0, or if it is a keras/tensorflow bug. In the later case, is there any workaround I can apply?
It's a bug that appeared in TF 2.10 and hasn't been solved as of TF 2.12.
For what is worth, EfficientNetV2s don't have this issue (tested on TF 2.12).