tensorflow: [TF2.0] KerasLayer cannot be loaded from .h5

Please make sure that this is a bug. As per our GitHub Policy, we only address code/doc bugs, performance issues, feature requests and build/installation issues on GitHub. tag:bug_template

System information Run on Colab example from copied from tensorflow/hub: https://colab.research.google.com/drive/1ymmzlWHieCuXR7UjXezULnYsNsDeHc9m

Version: 2.0.0-alpha0 Eager mode: True Hub version: 0.3.0 GPU is available

The issue When I try to run Colab example from tensorflow/hub for image retraining and then save/load model according to the documentation, there are some issues with it.

After retraining, when we try to save the model::

# Save the model
model.save('path_to_my_model.h5')

This error occurs (already described here: #26811 )

//...skipped, see ticket above for details ^
h5py/_objects.pyx in h5py._objects.with_phil.wrapper()
h5py/_objects.pyx in h5py._objects.with_phil.wrapper()
h5py/h5o.pyx in h5py.h5o.link()
RuntimeError: Unable to create link (name already exists)

But actually path_to_my_model.h5 file is created. Then when I try to load it via:

# Recreate the exact same model purely from the file
new_model = keras.models.load_model('path_to_my_model.h5')
new_model.summary()

I get an error:

ValueError: Unknown layer: KerasLayer

Entire stacktrace:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-15-0ceb2f1e6ad5> in <module>()
----> 1 new_model = keras.models.load_model('path_to_my_model.h5')
      2 
      3 new_model.summary()

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/saving/hdf5_format.py in load_model(filepath, custom_objects, compile)
    213     model_config = json.loads(model_config.decode('utf-8'))
    214     model = model_config_lib.model_from_config(model_config,
--> 215                                                custom_objects=custom_objects)
    216 
    217     # set weights

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/saving/model_config.py in model_from_config(config, custom_objects)
     53                     '`Sequential.from_config(config)`?')
     54   from tensorflow.python.keras.layers import deserialize  # pylint: disable=g-import-not-at-top
---> 55   return deserialize(config, custom_objects=custom_objects)
     56 
     57 

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/layers/serialization.py in deserialize(config, custom_objects)
     93       module_objects=globs,
     94       custom_objects=custom_objects,
---> 95       printable_module_name='layer')

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/utils/generic_utils.py in deserialize_keras_object(identifier, module_objects, custom_objects, printable_module_name)
    190             custom_objects=dict(
    191                 list(_GLOBAL_CUSTOM_OBJECTS.items()) +
--> 192                 list(custom_objects.items())))
    193       with CustomObjectScope(custom_objects):
    194         return cls.from_config(cls_config)

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/sequential.py in from_config(cls, config, custom_objects)
    349     for layer_config in layer_configs:
    350       layer = layer_module.deserialize(layer_config,
--> 351                                        custom_objects=custom_objects)
    352       model.add(layer)
    353     if not model.inputs and build_input_shape:

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/layers/serialization.py in deserialize(config, custom_objects)
     93       module_objects=globs,
     94       custom_objects=custom_objects,
---> 95       printable_module_name='layer')

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/utils/generic_utils.py in deserialize_keras_object(identifier, module_objects, custom_objects, printable_module_name)
    179     config = identifier
    180     (cls, cls_config) = class_and_config_for_serialized_keras_object(
--> 181         config, module_objects, custom_objects, printable_module_name)
    182 
    183     if hasattr(cls, 'from_config'):

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/utils/generic_utils.py in class_and_config_for_serialized_keras_object(config, module_objects, custom_objects, printable_module_name)
    164     cls = module_objects.get(class_name)
    165     if cls is None:
--> 166       raise ValueError('Unknown ' + printable_module_name + ': ' + class_name)
    167   return (cls, config['config'])
    168 

ValueError: Unknown layer: KerasLayer

Entire colab to review: https://colab.research.google.com/drive/1ymmzlWHieCuXR7UjXezULnYsNsDeHc9m