Hub: Unable to load a TF2-style SavedModel from TF Hub into a Keras model

[BrentDelarue]

Describe the bug

I tried loading the efficientnet/b0/feature-vector/1 TF2.0 savedmodel from TF Hub using tensorflow_hub.KerasLayer into a Keras Sequential model.
But wasn't able to load the savedmodel due to a version mismatch error.

System information

• OS Platform and Distribution: Debian GNU/Linux 9 (stretch)
• Tensorflow version: 2.0.0
• Tensorflow-hub version: 0.8.0
• Python version: 3.6.9

To Reproduce

import tensorflow_hub as hub
module = hub.KerasLayer("https://tfhub.dev/tensorflow/efficientnet/b0/feature-vector/1", trainable=True, input_shape=(224,224,3))

Logs

Traceback (most recent call last):
File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/framework/importer.py", line 501, in _import_graph_def_internal
graph._c_graph, serialized, options) # pylint: disable=protected-access
tensorflow.python.framework.errors_impl.InvalidArgumentError: NodeDef mentions attr 'exponential_avg_factor' not in Op y:T, batch_mean:U, batch_variance:U, reserve_space_1:U, reserve_space_2:U, reserve_space_3:U; attr=T:type,allowed=[DT_HALF, DT_BFLOAT16, DT_FLOAT]; attr=U:type,allowed=[DT_FLOAT]; attr=epsilon:float,default=0.0001; attr=data_format:string,default="NHWC",allowed=["NHWC", "NCHW"]; attr=is_training:bool,default=true>; NodeDef: {{node FusedBatchNormV3}}. (Check whether your GraphDef-interpreting binary is up to date with your GraphDef-generating binary.).

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
File "trainClassification.py", line 19, in
experimenter.run_training_experiment(config)
File "/hostdir/ai/BDA_TF2_test/BDA_tf2_framework/neuralnetwork/trainingexperimenter.py", line 49, in run_training_experiment
self.load_configuration(config_path)
File "/hostdir/ai/BDA_TF2_test/BDA_tf2_framework/neuralnetwork/trainingexperimenter.py", line 56, in load_configuration
self.apply_configuration()
File "/hostdir/ai/BDA_TF2_test/BDA_tf2_framework/neuralnetwork/trainingexperimenter.py", line 69, in apply_configuration
self.configure_trainer()
File "/hostdir/ai/BDA_TF2_test/BDA_tf2_framework/neuralnetwork/trainingexperimenter.py", line 90, in configure_trainer
self.model = self.network.get_network_model()
File "/hostdir/ai/BDA_TF2_test/BDA_tf2_framework/neuralnetwork/networkinterface.py", line 36, in get_network_model
self.construct_network()
File "/hostdir/ai/BDA_TF2_test/myfiles/mynetwork.py", line 52, in construct_network
module = hub.KerasLayer("https://tfhub.dev/tensorflow/efficientnet/b0/feature-vector/1", trainable=True, input_shape=(224,224,3))
File "/home/lumos/.local/lib/python3.6/site-packages/tensorflow_hub/keras_layer.py", line 146, in __init__
self._func = load_module(handle, tags)
File "/home/lumos/.local/lib/python3.6/site-packages/tensorflow_hub/keras_layer.py", line 398, in load_module
return module_v2.load(handle, tags=tags)
File "/home/lumos/.local/lib/python3.6/site-packages/tensorflow_hub/module_v2.py", line 102, in load
obj = tf_v1.saved_model.load_v2(module_path, tags=tags)
File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/saved_model/load.py", line 517, in load
return load_internal(export_dir, tags)
File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/saved_model/load.py", line 541, in load_internal
export_dir)
File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/saved_model/load.py", line 114, in __init__
meta_graph.graph_def.library))
File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/saved_model/function_deserialization.py", line 312, in load_function_def_library
copy, copy_functions=False)
File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/framework/function_def_to_graph.py", line 65, in function_def_to_graph
importer.import_graph_def_for_function(graph_def, name="")
File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/framework/importer.py", line 412, in import_graph_def_for_function
graph_def, validate_colocation_constraints=False, name=name)
File "/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/framework/importer.py", line 505, in _import_graph_def_internal
raise ValueError(str(e))
ValueError: NodeDef mentions attr 'exponential_avg_factor' not in Op y:T, batch_mean:U, batch_variance:U, reserve_space_1:U, reserve_space_2:U, reserve_space_3:U; attr=T:type,allowed=[DT_HALF, DT_BFLOAT16, DT_FLOAT]; attr=U:type,allowed=[DT_FLOAT];attr=epsilon:float,default=0.0001; attr=data_format:string,default="NHWC",allowed=["NHWC", "NCHW"]; attr=is_training:bool,default=true>; NodeDef: {{node FusedBatchNormV3}}. (Check whether your GraphDef-interpreting binary is up to date with yourGraphDef-generating binary.).

[gowthamkpr]

Was able to reproduce this bug.

[ClementWalter]

Having the same issue here with:

• tf 2.1.0
• Python 3.6.10 (default, Dec 19 2019, 23:04:32) [GCC 5.4.0 20160609] on linux

[arnoegw]

Please try TF2.2. It adds support for the new attr called out in the error message.

[BrentDelarue]

Please try TF2.2. It adds support for the new attr called out in the error message.

The issue is resolved when upgrading to Tensorflow version 2.2.0.

[arnoegw]

Thank you for the verification, and sorry for the confusion. We should at least update the docs to be clear on this.

[arnoegw]

Thanks again for the report. Docs have been updated (e.g., see https://tfhub.dev/tensorflow/efficientnet/b0/classification/1)