Tensorflow model import to Java

Question

I have been trying to import and make use of my trained model (Tensorflow, Python) in Java.

I was able to save the model in Python, but encountered problems when I try t

Answer 1

The Java importGraphDef() function is only importing the computational graph (written by tf.train.write_graph in your Python code), it isn't loading the values of trained variables (stored in the checkpoint), which is why you get an error complaining about uninitialized variables.

The TensorFlow SavedModel format on the other hand includes all information about a model (graph, checkpoint state, other metadata) and to use in Java you'd want to use SavedModelBundle.load to create session initialized with the trained variable values.

To export a model in this format from Python, you might want to take a look at a related question Deploy retrained inception SavedModel to google cloud ml engine

In your case, this should amount to something like the following in Python:

def save_model(session, input_tensor, output_tensor):
  signature = tf.saved_model.signature_def_utils.build_signature_def(
    inputs = {'input': tf.saved_model.utils.build_tensor_info(input_tensor)},
    outputs = {'output': tf.saved_model.utils.build_tensor_info(output_tensor)},
  )
  b = saved_model_builder.SavedModelBuilder('/tmp/model')
  b.add_meta_graph_and_variables(session,
                                 [tf.saved_model.tag_constants.SERVING],
                                 signature_def_map={tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature})
  b.save() 

And invoke that via save_model(session, x, yhat)

And then in Java load the model using:

try (SavedModelBundle b = SavedModelBundle.load("/tmp/mymodel", "serve")) {
  // b.session().run(...)
}

Hope that helps.