Hub: How to load hub modules in java

[haitong]

As title. Is there a way to recreate the tensorflow graph in java for a hub module?

[vbardiovskyg]

Hi @haitong,

there is currently no Java API for TensorFlow Hub. You could try exporting the module as a SavedModel (some hint here, although could be made much easier: https://github.com/tensorflow/hub/issues/99) and then loading as you would normally do in Java.

[haitong]

I am able to download one of the tensorflow hub module using "wget https://storage.googleapis.com/tfhub-modules/google/nnlm-en-dim128-with-normalization/1.tar.gz"

After extracting the module using tar -xzvf 1.tar.gz, I found that there are no tags for the module:
saved_model_cli show --dir . gives The given SavedModel contains the following tag-sets: with no tag-sets.

Without tags I am not able to load it using Java SavedModelBundle load API [public static SavedModelBundle load(String exportDir, String... tags)](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/java/src/main/java/org/tensorflow/SavedModelBundle.java#L94) as it explicitly expects tags.

Is it expected that there is no tags info for tensorflow hub modules?

[vbardiovskyg]

Yes, the hub module cannot be directly loaded as a SavedModel and it has to be first exported as one. The issue I linked above shows how to do it. In short you declare in a clean graph a placeholder, then apply tensorflow hub module on the placeholder, then export the SavedModel using tf.saved_model.simple_save.

[haitong]

Thanks for the info! Really appreciate it!

[alexandergunnarson]

For other people's benefit — after scouring the internet, here's how I was able to do this:

# Run this via `python3`

import tensorflow as tf
import tensorflow_hub as hub

def save_module(url, save_path, input_type):
    with tf.Graph().as_default():
        module = hub.Module(url)
        model_input = tf.placeholder(input_type, name="input")
        model_output = tf.identity(module(model_input), name="output")
        with tf.Session() as session:
            session.run(tf.global_variables_initializer())
            tf.saved_model.simple_save(
                session,
                save_path,
                inputs={'input': model_input},
                outputs={'output': model_output},
                legacy_init_op=tf.initializers.tables_initializer(name='init_all_tables'))

# save_module("https://tfhub.dev/google/universal-sentence-encoder-large/3", "./saved-module", tf.string)

And then in Java:

import org.tensorflow.SavedModelBundle;
import org.tensorflow.Tensor;

SavedModelBundle savedModelBundle = SavedModelBundle.load("./saved-module", "serve");

Tensor result = 
  savedModelBundle
  .session()
  .runner()
  // Note that I implemented `convert1DArrayTo2DArray` in Clojure, so I didn't paste here
  .feed("input", convert1DArrayTo2DArray("Hello world!".getBytes("UTF-8")))
  .fetch("output")
  .run()
  .get(0);

[cowwoc]

@vbardiovskyg has the situation changed? Is it possible to load hub modules in Java using TensorFlow 2?

[vbardiovskyg]

Hi @cowwoc, I don't think there has been any change in SavedModel's Java APIs recently.

You can load the hub modules as in the example above.

As for the new TF2.0 SavedModels that are also hosted on tfhub.dev, they will be loadable directly if they were exported with a serving signature.

[bfurlan]

code for TF 2 models (based on @alexandergunnarson solution):

import tensorflow.compat.v1 as tf
import tensorflow_hub as hub
tf.disable_v2_behavior() 

def save_module(url, save_path, input_type):
    with tf.Graph().as_default():
        module = hub.KerasLayer(url)
        model_input = tf.placeholder(input_type, name="input")
        model_output = tf.identity(module(model_input), name="output")
        with tf.Session() as session:
            session.run(tf.global_variables_initializer())
            tf.saved_model.simple_save(
                session,
                save_path,
                inputs={'input': model_input},
                outputs={'output': model_output},
                legacy_init_op=tf.initializers.tables_initializer(name='init_all_tables'))

# save_module("https://tfhub.dev/google/universal-sentence-encoder-large/5", "models/tmp/universal-sentence-encoder", tf.string)

and in Java:

import org.tensorflow.SavedModelBundle;
import org.tensorflow.Tensor;
import java.nio.charset.StandardCharsets;
import java.io.UnsupportedEncodingException;

    SavedModelBundle savedModelBundle = SavedModelBundle.load(path, "serve");

    float[][] embed(String[] values) throws UnsupportedEncodingException{

        byte[][] input = new byte[values.length][];
        for (int i = 0; i < values.length; i++) {
            String val = values[i];
            input[i] = val.getBytes(StandardCharsets.UTF_8);

        }

        Tensor t = Tensor.create(input);

        Tensor result =
                savedModelBundle
                        .session()
                        .runner()
                        .feed("input", t)
                        .fetch("output")
                        .run()
                        .get(0);

        float[][] output = new float[values.length][512];
        result.copyTo(output);

        return output;

    }

[samikrc]

Is there a python code piece where we can save the hub module using TF2 (instead of disabling the TF1 behavior of using TF Session etc), and load it in JVM using TF2? Or it is necessary to stick to the TF1 paradigm? Thanks.

[cowwoc]

@samikrc You don't need to save the hub module using python and re-loading it using the JVM. You can load it directly from the JVM using SavedModelBundle.load(). See https://stackoverflow.com/a/61968561/14731