Search code examples
pythontensorflowtensorflow-hub

TensorFlow Hub error when Saving model as H5 or SavedModel


I want to use this TF Hub asset: https://tfhub.dev/google/imagenet/resnet_v1_50/feature_vector/3

Versions:

Version:  1.15.0-dev20190726
Eager mode:  False
Hub version:  0.5.0
GPU is available

Code

feature_extractor_url = "https://tfhub.dev/google/imagenet/resnet_v1_50/feature_vector/3"
feature_extractor_layer = hub.KerasLayer(module,
                                         input_shape=(HEIGHT, WIDTH, CHANNELS))

I get:

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.

I tried:

module = hub.Module("https://tfhub.dev/google/imagenet/resnet_v1_50/feature_vector/3",
                    tags={"train"})
feature_extractor_layer = hub.KerasLayer(module, 
                                         input_shape=(HEIGHT, WIDTH, CHANNELS))

But when I try to save the model I get:

tf.keras.experimental.export_saved_model(model, tf_model_path)
# model.save(h5_model_path) # Same error 

NotImplementedError: Can only generate a valid config for `hub.KerasLayer(handle, ...)`that uses a string `handle`.
Got `type(handle)`: <class 'tensorflow_hub.module.Module'>

Tutorial here


Solution

  • It's been a while, but assuming you have migrated to the TF2, this can easily be accomplished with the most recent model version as follows:

    import tensorflow as tf
    import tensorflow_hub as hub
    
    num_classes=10 # For example
    m = tf.keras.Sequential([
        hub.KerasLayer("https://tfhub.dev/google/imagenet/resnet_v1_50/feature_vector/5", trainable=True)
        tf.keras.layers.Dense(num_classes, activation='softmax')
    ])
    m.build([None, 224, 224, 3])  # Batch input shape.
    
    # train as needed
    
    m.save("/some/output/path")
    

    Please update this question if that doesn't work for you. I believe your issue arose from mixing hub.Module with hub.KerasLayer. The model version you were using was in TF1 Hub format, so within TF1 it is meant to be used exclusively with hub.Module, and not mixed with hub.KerasLayer. Within TF2, hub.KerasLayer can load TF1 Hub format models directly from their URL for composition in larger models, but they cannot be fine-tuned.

    Please refer to this compatibility guide for more information