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: #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

[Udayk02]

@cowwoc @haitong The same code is not working right now. Are there any changes that you have done?

[cowwoc]

Sorry, I migrated to pytorch so I have no idea :(

[ip2016]

If someone still needs solution, here is an update.
Turns out for TF 2.0 models it actually works out of the box, so it doesn't need to be saved in python. Just download model directly from the hub (i.e. https://tfhub.dev/google/universal-sentence-encoder/4)

In order to load model in java you would need to know input and output names. To get it you can run "saved_model_cli" utility available from python tensorflow installation (i.e. saved_model_cli show --dir <model_folder> --all). The output will be similar to following one:

signature_def['serving_default']:
  The given SavedModel SignatureDef contains the following input(s):
    inputs['inputs'] tensor_info:
        dtype: DT_STRING
        shape: unknown_rank
        name: serving_default_inputs:0
  The given SavedModel SignatureDef contains the following output(s):
    outputs['output_0'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, 512)
        name: StatefulPartitionedCall:0
  Method name is: tensorflow/serving/predict

Here input name is serving_default_inputs and output name is StatefulPartitionedCall
So now you can call the model in java as follow:

SaveModelBundle bundle = SaveModelBundle.load("<model_folder>", "serve")
String[] messages = [.....]
NdArray<String> ar = NdArrays.ofObjects(String.class, StdArrays.shapeOf(messages))
Tensor tensor = TString.tensorOf(ar)

TFloat32 result = bundle.session().runner().feed("serving_default_inputs", tensor).fetch("StatefulPartitionedCall").run().get(0)

To get inputs/outputs in Java, you can use:
String input = bundle.metaGraphDef().getSignatureDefMap().get("serving_default").getInputsMap()["inputs"].getName()
Similar for output.