Hub: Can't load EfficientNet using TensorFlow 2

[artemmavrin]

I am unable to load EfficientNet models using

• TensorFlow: 2.0.0
• TensorFlow Hub: 0.6.0

For example, I tried the code stub at the efficientnet/b0/feature-vector page:

import tensorflow_hub as hub
module = hub.Module("https://tfhub.dev/google/efficientnet/b0/feature-vector/1")

Error:

RuntimeError: Exporting/importing meta graphs is not supported when eager execution is enabled. No graph exists when eager execution is enabled.

If I try to load the model using hub.KerasLayer instead, I get a different error.

import tensorflow_hub as hub
layer = hub.KerasLayer("https://tfhub.dev/google/efficientnet/b0/feature-vector/1")

Error:

ValueError: Importing a SavedModel with tf.saved_model.load requires a 'tags=' argument if there is more than one MetaGraph. Got 'tags=None', but there are 2 MetaGraphs in the SavedModel with tag sets [['train'], []]. Pass a 'tags=' argument to load this SavedModel.

---

The KerasLayer error might be fixed by modifying hub/tensorflow_hub/keras_layer.py by adding a tags argument to KerasLayer.__init__, so that

https://github.com/tensorflow/hub/blob/9bc54d7920e2b3faa63eada0e2b6eee70cc07a79/tensorflow_hub/keras_layer.py#L98

becomes

   def __init__(self, handle, trainable=False, arguments=None, tags=None, **kwargs):

and modifying the line

https://github.com/tensorflow/hub/blob/9bc54d7920e2b3faa63eada0e2b6eee70cc07a79/tensorflow_hub/keras_layer.py#L109

to

      self._func = module_v2.load(handle, tags=tags)

Then I could pass, for example, tags=['train'] to load the EfficientNet model as a KerasLayer, as in

import tensorflow_hub as hub
layer = hub.KerasLayer("https://tfhub.dev/google/efficientnet/b0/feature-vector/1",
                       tags=['train'])

[rmothukuru]

Could reproduce the Error with Tensorflow Version 2.0 and TF Hub version 0.6. Here is the Github Gist. Thanks!

[maximneumann]

This EfficientNet module is an old SavedModel format, that can be loaded in TF2 for example with:

module = hub.load("https://tfhub.dev/google/efficientnet/b0/feature-vector/1", tags=[])

This is briefly documented in https://www.tensorflow.org/hub/migration_tf2#loading_old_hubmodules

To load it as a KerasLayer, one could write a wrapper, or making the change you suggested. The hub team is currently working on a fix for this (probably to be released with hub v0.7).

[maximneumann]

Here is an example KerasLayer wrapper for the old SavedModels, if you need it right now:

class Wrapper(tf.train.Checkpoint):
  def __init__(self, spec):
    super(Wrapper, self).__init__()
    self.module = hub.load(spec, tags=[])
    self.variables = self.module.variables
    self.trainable_variables = []
  def __call__(self, x):
    return self.module.signatures["default"](x)["default"]

m = hub.KerasLayer(Wrapper("https://tfhub.dev/google/efficientnet/b0/feature-vector/1"))

Note that it's only a temporal fix and will probably not be required with hub v0.7.