How to load tensorflow .pb model in java exported from keras

user_7128294 0
Viewed 2233

I'm trying to load a model in java which is originally saved in keras in Java so that I can do inference in-process in an existing production system which runs in Java.

I didn't see a way to load Keras h5 models easily in Java, so I'm attempting to first convert it into a .pb file by using simple_save, and then loading it using the default tag for simple_save. I tried saving the graph directly using a freeze_session routine and tf.train.write_graph, but I had the same error.

Here's the code to save my model to a .pb file

# my model has two input tensors and one output tensor
inputs = {'input_1': model.inputs[0], 'input_2' : model.inputs[1]}
outputs = {'output_1' : model.outputs[0]}

tf.saved_model.simple_save(K.get_session(), 'output_dir', inputs=inputs, outputs=outputs)

Here's my Java code for loading the model, using the default tag for saved_model:

SavedModelBundle model = SavedModelBundle.load("output_dir", "serve");

This is resulting in the error:

Exception in thread "main" org.tensorflow.TensorFlowException: Could not find SavedModel .pb or .pbtxt at supplied export directory path: output_dir

Any idea what I might be doing wrong? I know that simple_save is deprecated, but I'm just trying to get anything to work at this point.

keras tensorflow java
2 Answers
user_131236390

I looked at the native source code which loads the model, and it turns out there is a hard-coded file name "saved_model.pb", or for the text version "saved_model.pbtxt" it expects in the directory (which wasn’t specified in the documentation I looked at).

user_71282940
Related