Python中的Tensorflow Java Api`toGraphDef`等价是什么？

Question

I am using the Tensorflow Java Api to load an already created Tensorflow model into the JVM. I am using this as an example: tensorflow/examples/LabelImage.java

Here is my simple scala code:

import java.nio.file.{Files, Path, Paths}
import org.tensorflow.{Graph, Session, Tensor}

def readAllBytesOrExit(path: Path): Array[Byte] = Files.readAllBytes(path)
val graphDef = readAllBytesOrExit(Paths.get("PATH_TO_A_SINGLE_FILE_DESCRIBING_TF_MODEL.pb"))
val g = new Graph()
g.importGraphDef(graphDef)
val session = new Session(g)
val result: Tensor = session.runner().feed("input", image).fetch("output").run().get(0))

How do I save my model to get both the Session and the Graph stored in the same file. as described in the "PATH_TO_A_SINGLE_FILE_DESCRIBING_TF_MODEL.pb" above.

Described here it mentions:

The serialized representation of the graph, often referred to as a GraphDef, can be generated by toGraphDef() and equivalents in other language APIs.

What are the equivalents in other language APIs? I dont find it obvious

Note: I already looked at the mnist_saved_model.py under tensorflow_serving but saving it through that procedure gives me a .pb file and a variables folder. When trying to load that .pb file I get: java.lang.IllegalArgumentException: Invalid GraphDef

Answer 1

Currently with the Java API of tensorflow, I only found how to save a graph as a graphDef (ie without its variables and meta-data). This can be done by just writing the Array[Byte] to a file:

Files.write(Paths.get(modelDir, modelName), myGraph.toGraphDef)

Here myGraph is a java object from the Graph class .

I would suggest to save your model from the Python API, using the SavedModel api defined here. It will save your model in a folder with both the serialized graph in a .pb file and the variables in a folder. Note the tag_constants you use as you'll need it in your scala/java code to load the model with the variables. Then the graph and session with variables are easily loaded with the SavedModelBundle java class from the java api. It returns you a wrapper with both the graph and the session containing the variables values:

val model = SavedModelBundle.load(modelDir, modelTag)

If you already tried this, maybe you can share your code to see why it returned an invalid GraphDef.

Another option is to freeze your graph, ie you turned your variable nodes into constant Nodes so everything is self-contained in the .pb file. Mores infos here for the freezing part