CODEWITHSUNDEEP
javatensorflowkeras
Nathaniel Powell
Nathaniel Powell

Reputation: 265

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

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.

Upvotes: 2

Views: 3160

Answers (2)

Frank Liu
Frank Liu

Reputation: 336

Now you can use Deep Java Library (DJL) to load Keras model in Java and run inference. DJL internally use tensorflow java and provide high level API to make it easy to run inference and trainging. Checkout the github repo: https://github.com/awslabs/djl

There is a blogpost: https://towardsdatascience.com/detecting-pneumonia-from-chest-x-ray-images-e02bcf705dd6

And the demo project can be found: https://github.com/aws-samples/djl-demo/blob/master/pneumonia-detection/README.md

Upvotes: 2

Nathaniel Powell
Nathaniel Powell

Reputation: 265

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).

Upvotes: 1

Related Questions