Refer resource folder and not the content in jar

Viewed 301

We have a simple ML model, compiled and saved as SavedModel/*.pb format. We load the SavedModel/ using TensorFlow 1.5 (Java) for inferencing.

We are using the below approach to load the model:

String path = 'models/SavedModel'
File modelFile = new File(getClass().getClassLoader().getResource(path).getPath());

model = SavedModelBundle.loader(modelFile.getAbsolutePath())
    .withTags("serve")
    .load();
Graph g = model.graph();
...

We are able to inference the model and get the output using IDE but it does not work, once we build a Jar.

Important notes:

  1. We have kept the SavedModel in resources dir, like src/main/resources/models/SavedModel/*.pb.
  2. SavedModelBundle.loader takes String exportDir as first parameter. exportDir is the directory path containing a saved model.
  3. Once we create the Jar of the project, we are able to see models/SavedModel/* on the root of the Jar.

We are not able to refer the correct path of the exportDir within the jar. Can anyone please help me here? I'm new to Java world!

1 Answers

Unfortunately, you cannot do this, at least not that easily.

TensorFlow does not read the saved model directory from the JVM but from its native C++ library. This library understand regular file paths but not Java resource paths. The reason why it works from the IDE is because most IDEs work directly with the classes compiled and stored on the file system, they don't have to deal with archives as the command line does.

For example, if I ran this little snippet of code:

public class Main {

  public static void main(String[] args) throws Exception {
    String path = Main.class.getResource("model").getPath();
    System.out.println("Loading model at " + path);
    SavedModelBundle.load(path, "serve");
  }
}

From the IDE, I get:

Loading model at /Users/klessard/Documents/Projects/MachineLearning/Sources/quick-java-test/target/classes/model
2020-05-17 01:53:31.355176: I tensorflow/cc/saved_model/reader.cc:31] Reading SavedModel from: /Users/klessard/Documents/Projects/MachineLearning/Sources/quick-java-test/target/classes/model
2020-05-17 01:53:31.428273: I tensorflow/cc/saved_model/reader.cc:54] Reading meta graph with tags { serve }
2020-05-17 01:53:31.787746: I tensorflow/cc/saved_model/loader.cc:182] Restoring SavedModel bundle.
2020-05-17 01:53:33.762491: I tensorflow/cc/saved_model/loader.cc:132] Running initialization op on SavedModel bundle.
2020-05-17 01:53:34.178367: I tensorflow/cc/saved_model/loader.cc:285] SavedModel load for tags { serve }; Status: success. Took 2823199 microseconds.

but running from a JAR with the following command line java -jar target/quick-java-test-1.0-SNAPSHOT-jar-with-dependencies.jar will give:

Loading model at file:/Users/klessard/Documents/Projects/MachineLearning/Sources/quick-java-test/target/quick-java-test-1.0-SNAPSHOT-jar-with-dependencies.jar!/model
2020-05-17 01:55:25.701852: I tensorflow/cc/saved_model/reader.cc:31] Reading SavedModel from: file:/Users/klessard/Documents/Projects/MachineLearning/Sources/quick-java-test/target/quick-java-test-1.0-SNAPSHOT-jar-with-dependencies.jar!/model
2020-05-17 01:55:25.701937: I tensorflow/cc/saved_model/loader.cc:285] SavedModel load for tags { serve }; Status: fail. Took 101 microseconds.
Exception in thread "main" org.tensorflow.TensorFlowException: Could not find SavedModel .pb or .pbtxt at supplied export directory path: file:/Users/klessard/Documents/Projects/MachineLearning/Sources/quick-java-test/target/quick-java-test-1.0-SNAPSHOT-jar-with-dependencies.jar!/model
        at org.tensorflow.SavedModelBundle.load(Native Method)
        at org.tensorflow.SavedModelBundle.access$000(SavedModelBundle.java:27)
        at org.tensorflow.SavedModelBundle$Loader.load(SavedModelBundle.java:32)
        at org.tensorflow.SavedModelBundle.load(SavedModelBundle.java:95)
        at Main.main(Main.java:10)

Compare both paths: From the IDE, it points to the target/classes folder on the file system (which is the default output folder for classes when compiling with Maven), while from the command line, it points to a directory inside a JAR (quick-java-test-1.0-SNAPSHOT-jar-with-dependencies.jar!/model), which is specific to Java. The C++ native library has no clue how to resolve this path and fail loading the model.

So I suggest that you store your saved model on the file system and use regular file paths to load it (if you use Docker, that could be done when building the image).

Another solution also is to import the graph using the legacy format, as it is passed directly as a serialized proto message instead of reading from a directory. But I'm not sure how long this format will still be supported by TensorFlow.

Finally, you could also have your saved model directory archived in a zip file under your resources, that you unzip to a tmp folder in Java before loading the model using SavedModelBundle.load.

Related