Hub: Feature request: using a hub.Module within tf.dataset.map

[akshitj1]

Hi, I am trying to apply module to batch of images, But I am unable to figure out how to apply module as map function for each batch with it's own graph.

This is how my function looks

def get_image_embeddings(image_files_list):
    with tf.Graph().as_default():
        config = tf.ConfigProto()
        config.gpu_options.allow_growth=True
        with tf.Session(config=config) as sess:
            NASNET_MOBILE='https://tfhub.dev/google/imagenet/nasnet_mobile/feature_vector/1'
            module = hub.Module(NASNET_MOBILE)

            sess.run(tf.global_variables_initializer())
            sess.run(tf.tables_initializer())

            image_files = tf.convert_to_tensor(image_files_list, dtype=tf.string)
            image_files_dataset = tf.data.Dataset.from_tensor_slices(image_files)
            expected_width, expected_height= hub.get_expected_image_size(module)
            images=image_files_dataset.map(lambda image: preprocess_image(image, expected_width, expected_height))
            image_batches=images.batch(100)
            embs = image_batches.map(module)
            embs = sess.run(embs)
            sess.close()
            return embs

but when I try to execute this, I am getting.
RuntimeError: Module must be applied in the graph it was instantiated for.

I took a look at this issue: #54
but couldn't relate solution to my problem.
Any help would be appreciated.

[arnoegw]

Hi akshitj1! Thanks for the report, and sorry for the inconvenience! Turns out a hub.Module cannot be used with Dataset.map(). We want to fix that. I'll keep this issue open for reference.

In the meantime, would it work for your case to work around this with

images_batches = ... # As before.
embs = module(embs)
return sess.eval()  # Returning out of  with tf.Session...  closes session.

?

[akshitj1]

Hi @arnoegw
I went with the dataset api because plain tensor batches was not working for me. Below code just keeps stuck after this statement. there is no debug/error/warn message after this statement.
INFO:tensorflow:Saver not created because there are no variables in the graph to restore
below is the code I had used:

def get_image_embeddings(image_files_list):
    with tf.Graph().as_default():
        with tf.Session() as sess:
            # load and batch images
            image_files = tf.convert_to_tensor(image_files_list, dtype=tf.string)
            images = tf.map_fn(preprocess_image, image_files, back_prop=False, dtype=tf.float32)
            image_batches = tf.train.batch([images], batch_size=10)

            # load module
            NASNET_MOBILE='https://tfhub.dev/google/imagenet/nasnet_mobile/feature_vector/1'
            module = hub.Module(NASNET_MOBILE)

            sess.run(tf.global_variables_initializer())
            sess.run(tf.tables_initializer())


            # find embeddings for each batches
            embs_batches = tf.map_fn(module, image_batches, back_prop=False)
            return embs_batches.eval()

I have a feeling I am doing some obvious mistake.

[arnoegw]

Hmm, I can't see right now what exactly goes wrong with that code, but it shouldn't be necessary to use tf.map_fn in the first place: the image modules from tfhub.dev all accept a batch of input images. See https://www.tensorflow.org/hub/common_signatures/images#image_input

[akshitj1]

Thanks @arnoegw . That worked.

[arnoegw]

Glad to hear that the workaround works!

Keeping this open for the original request, that is, Dataset.map(module).

[jannea]

Any updates on using hub.Module inside tf.dataset.map ?

[Harshini-Gadige]

@arnoegw Any updates on this enhancement ?

[arnoegw]

For TensorFlow 2.0, TF Hub will ship modules that are native TF2 SavedModels containing @tf.functions. Those "should just work" in Dataset.map(). To get a preview of the upcoming Hub image modules, see examples/colab/tf2_image_retraining.ipynb.

We don't plan to go back and fix that for TF1.x, so I'm closing this issue.

[omoindrot]

@arnoegw : I've tried using tf.data.Dataset.map with the new TF Hub modules and it works ! 🍾

One issue is that the tf.data.Dataset.map operation only supports CPU, is there a plan to add GPU support in the future?

The error is:

NotFoundError: No registered 'MapDataset' OpKernel for GPU devices compatible with node {{node MapDataset}}

[arnoegw]

@omoindrot, thanks for confirming.

Regarding GPU support for tf.data.Dataset.map: that's an excellent question for https://github.com/tensorflow/tensorflow (not /hub). It sounds about right that the MapDataset op itself executes on CPU, but where does it run the function inside?