:images-asciidoc: https://raw.githubusercontent.com/spring-cloud/stream-applications/master/functions/common/tensorflow-common/src/main/resources/images/
= Programming Model for TensorFlow Inference
Programming model builds on the https://docs.oracle.com/javase/8/docs/api/java/util/function/package-summary.html[Java Function API], the `TF Java Ops API` and few basic data structure that together help to unify and streamline the building of TensorFlow inference pipelines.
Quick Start: just add the following dependency:
[source,XML]
----
<dependency>
<groupId>org.springframework.cloud.fn</groupId>
<artifactId>tensorflow-common</artifactId>
<version>${revision}</version>
</dependency>
----
== Programming Model
Implementing a real-time TensorFlow Inference in Java, typically leverages the TF Java API for loading and scoring the pre-trained models. But the logic used to convert the upstream data into model input Tensors (e.g. pre-processor) and in turn to convert the inferred Tensors back into application data (e.g. post-processor) is commonly implemented in plain Java:
image::{images-asciidoc}/programming_model.png[TF Architecture, scaledwidth="70%"]
The Pre and Post processing steps could become very complex (check https://github.com/ildoonet/tf-pose-estimation[Pose Estimation] or https://github.com/davidsandberg/facenet[Face Recognition]) and computationally intensive. E.g. tons of math and image operations that are better fit for optimized TF utilities rather the plain Java math or AWT/2D/Canvas such.
Additionally, the unnecessary shuffling of data between the JVM and the native TF impacts the overall performance of the flow. The issue is apparent when multiple pre-trained TF models are combined (composed) or the same TF models are evaluated iteratively. In those cases the processed data is repeatedly being moved between the JVM Heap and the TF native memory.
The `Java Ops API` exposes the native https://www.tensorflow.org/versions/r1.9/api_docs/cc?hl=en[TF C++ Core API] offering comprehensive, native math, image, io and alike utilities. Later is useful for implementing the computational intensive logic by using the same TF tools and infrastructure used for running the pre-trained models.
Proposed programming model builds on the https://docs.oracle.com/javase/8/docs/api/java/util/function/package-summary.html[Java Function API], the Java Ops API and a basic data structure that together help to unify and streamline the building of TensorFlow inference pipelines. While focused on model-inference, this programming model is likely to be useful for building model-training pipelines as well.
The programming model leverages the following functional definition:
[source,Java]
----
Function<Map<String, Tensor<?>>, Map<String, Tensor<?>>>
----
and the corresponding method expression
[source,Java]
----
Map<String, Tensor<?>> apply( Map<String, Tensor<?>> feeds)
----
This function receives a map of named Tensors as an input and in turn returns a map of named Tensors. Because the input and the output formats are equivalent, functions with this signature can be reused and composed into larger, complex functions.
The names used in the inputs, and the outputs maps are strings of the form `operation_name: output_index` (output_index defaults to 0). The names must match the indexed operations inside the underlying TF graph. +
The https://www.tensorflow.org/api_docs/java/reference/org/tensorflow/Tensor[Tensor] class is a reference to the data used natively in the TF engine. The referenced data is not moved to JVM Heap unless explicitly materialized with `Tensor.copyTo()`. Exchanging Tensor references between functions prevents unnecessarily copying the data to the JVM heap. +
Proposed data structure fits well with existing https://www.tensorflow.org/api_docs/java/reference/org/tensorflow/Session.Runner[Session.Runner API], which accepts https://www.tensorflow.org/api_docs/java/reference/org/tensorflow/Session.Runner.html#feed(java.lang.String,%20org.tensorflow.Tensor%3C?%3E)[indexed operations] as an input feed and returns list of tensors predefined by https://www.tensorflow.org/api_docs/java/reference/org/tensorflow/Session.Runner.html#fetch(java.lang.String)[fetch indexed operations].
The https://github.com/spring-cloud/stream-applications/blob/master/functions/common/tensorflow-common/src/main/java/org/springframework/cloud/fn/common/tensorflow/GraphRunner.java[GraphRunner] and the https://github.com/spring-cloud/stream-applications/blob/master/functions/common/tensorflow-common/src/main/java/org/springframework/cloud/fn/common/tensorflow/GraphDefinition.java[GraphDefinition] are the core abstractions used to define, load and inference TensorFlow models. https://github.com/spring-cloud/stream-applications/blob/master/functions/common/tensorflow-common/src/main/java/org/springframework/cloud/fn/common/tensorflow/GraphRunner.java[GraphRunner] implements the Function (e.g. Fn<Map<S,T>, Map<S,T>>) definition and uses the TF Java API to run the underlying TF graph. The input Tensor map is fed to the Session Runner. After the graph is evaluated, a list of predefined fetch names is used to retrieve selected Tensors from the result as a named Tensor map. The https://github.com/spring-cloud/stream-applications/blob/master/functions/common/tensorflow-common/src/main/java/org/springframework/cloud/fn/common/tensorflow/GraphRunner.java#L70[withGraphDefinition(GraphDefinition)] method defines a new or loads a pre-trained TF graph, while the https://github.com/spring-cloud/stream-applications/blob/master/functions/common/tensorflow-common/src/main/java/org/springframework/cloud/fn/common/tensorflow/GraphRunner.java#L84[withSavedModel(path)] method helps to load a Tensorflow SavedModel.
The GraphDefinition argument is a functional interface and can therefore be used as the assignment target for a lambda expression or method reference.
Following snippets illustrates how to use the `withGraphDefinition` to define a new TF Graph that computes the `y1 = x1 * 2` expression:
[source,Java]
----
myGraph = new GraphRunner("x1", "y1")
.withGraphDefinition( tf ->
tf.withName("y1").math.mul(
tf.withName("x1").placeholder(Integer.class),
tf.constant(2)));
----
The x1 and y1 constructor arguments define the input and output Tensor names (technically indexed operation names) used to feed in and fetch out data to and from the defined model.
The GraphRunner https://github.com/spring-cloud/stream-applications/blob/master/functions/common/tensorflow-common/src/main/java/org/springframework/cloud/fn/common/tensorflow/AbstractGraphRunner.java#L65[apply] method helps evaluate/inference the so defined graph:
[source,Java]
----
Map<String, Tensor<?>> input = Collections.singletonMap("x1", Tensor.create(666));
result = myGraph.apply(input);
----
Similarly, we can load a frozen/pre-trained model (https://github.com/tensorflow/models/tree/master/research/slim/nets/mobilenet#pretrained-models[MobileNetV2] model in this case) from an archive using the https://github.com/spring-cloud/stream-applications/blob/master/functions/common/tensorflow-common/src/main/java/org/springframework/cloud/fn/common/tensorflow/ProtoBufGraphDefinition.java[ProtoBufGraphDefinitions] helper class.
[source,Java]
----
mibileNetV2 = new GraphRunner("input", "MobilenetV2/Predictions/Reshape_1")
.withGraphDefinition(
new ProtoBufGraphDefinition(
"https://storage.googleapis.com/mobilenet_v2/checkpoints/mobilenet_v2_1.4_224.tgz#mobilenet_v2_1.4_224_frozen.pb",
cacheModel));
----
You can load archives from `http://`, `file://` or `classpath://` locations.
For loading a `SavedModel` use the `GraphRunner#withSavedModel method like this:
[source,Java]
----
ssdMibileNetV1Coco = new GraphRunner( Arrays.asList("image_tensor"),
Arrays.asList("detection_boxes", “detection_scores”, “detection_classes”))
.withSavedModel( ”./ssd_mobilenet_v1_coco_2017_11_17/saved_model”, "serve");
----
An example inference pipeline based on the proposed Functional Programming Model would look similar to this:
image::{images-asciidoc}/tf_pipeline.png[TF pipelinne, scaledwidth="70%"]
Every GraphRunner instance in the pipeline uses either an in-place defined, or a pre-trained TF graph.
In practice, it would still be required to implement some input and output adapters for the logic that cannot be (or are not feasible to be) implemented with the native Java Ops API. But you have the freedom to choose what part of the processing logic to run as natively (e.g. Java Ops API) code and what is a plain Java.
Furthermore, we are not limited to GraphRunner but any custom https://docs.oracle.com/javase/8/docs/api/java/util/function/Function.html[Function]<Map<String, Tensor>, Map<String, Tensor>> implementations can be used in the processing pipelines. In fact the https://github.com/spring-cloud/stream-applications/blob/master/functions/common/tensorflow-common/src/main/java/org/springframework/cloud/fn/common/tensorflow/Functions.java[Functions] utilities use this approach.
When appropriate any custom Function, https://docs.oracle.com/javase/8/docs/api/java/util/function/Supplier.html[Supplier]<Map<String, Tensor>>, https://docs.oracle.com/javase/8/docs/api/java/util/function/Consumer.html[Consumer]<Map<String, Tensor>> or the rest of the https://docs.oracle.com/javase/8/docs/api/java/util/function/package-frame.html[java.util.function] classes and interfaces can be used.
For real time examples check the https://github.com/spring-cloud/stream-applications/tree/master/functions/function/image-recognition-function[image-recognition] and https://github.com/spring-cloud/stream-applications/tree/master/functions/function/semantic-segmentation-function[semantic-segmentation] implementations.
== Features
Following paragraphs discusses some features and techniques useful for composing graphs, memorizing and reusing intermediate Tensor values, managing the Tensor resources and so on.
=== Input and Output Contracts
The GraphRunner constructor expects two compulsory list fields: `feedNames` - list of names that the graph accepts as input Tensors and `fetchNames` - list of (Tensor) names that the graph would return.
[source,Java]
----
public GraphRunner(List<String> feedNames, List<String> fetchedNames)
----
Together those two lists define the input and output contract of the graph runner. +
The names used in the inputs, and the outputs maps are strings of the form `operation_name : output_index` (output_index defaults to 0). The names must match the indexed operations inside the underlying TF graph.
=== Composition
Because the https://github.com/spring-cloud/stream-applications/blob/master/functions/common/tensorflow-common/src/main/java/org/springframework/cloud/fn/common/tensorflow/GraphRunner.java[GraphRunner] function signature uses the same type for input and output parameters, the https://docs.oracle.com/javase/8/docs/api/java/util/function/Function.html[Functional] interface allows us to compose multiple graphs https://github.com/spring-cloud/stream-applications/blob/master/functions/common/tensorflow-common/src/main/java/org/springframework/cloud/fn/common/tensorflow/GraphRunner.java[GraphRunner] functions into a larger composite function:
[source,Java]
----
composed-graph = graph1.andThen(graph2)....andThen(graphN)
----
For example let's take two simple graphs: `G1 (y1 = x1 * 2)` and `G2 (y2 = x2 + 20)`. The composed graph `G = G1.andThen(G2)` is equivalent to `y = (x * 2 ) + 20`.
The https://github.com/spring-cloud/stream-applications/blob/master/functions/common/tensorflow-common/src/test/java/org/springframework/cloud/fn/common/tensorflow/FunctionComposition.java[FunctionComposition example] demonstrates how this works:
[source,Java]
----
try (
GraphRunner graph1 = new GraphRunner("x1", "y1")
.withGraphDefinition(tf -> tf.withName("y1").math.mul(
tf.withName("x1").placeholder(Integer.class),
tf.constant(2)));
GraphRunner graph2 = new GraphRunner("x2", "y2")
.withGraphDefinition(tf -> tf.withName("y2").math.add(
tf.withName("x2").placeholder(Integer.class),
tf.constant(20)));
Tensor x = Tensor.create(10);
) {
Map<String, Tensor<?>> result =
graph1.andThen(graph2).apply(Collections.singletonMap("x", x));
System.out.println("Result is: " + result.get("y2").intValue()); // Result is: 40
}
----
Note that the GraphRunner https://github.com/spring-cloud/stream-applications/blob/master/functions/common/tensorflow-common/src/main/java/org/springframework/cloud/fn/common/tensorflow/AbstractGraphRunner.java#L65[automatically binds] the singleton outputs (e.g fetch) with the singleton input (e.g. feeds). In the example above the GraphRunner automatically binds the `y1` tensor produced by `graph1` to the `x2` input placeholders expected by `graph2`.
==== Multiple inputs/outputs
When the composed graphs use multiple input and output parameters we need to explicitly bind the outputs from the upstream graph to the inputs of the downstream one.
For example let’s Graph1 produces two outputs (e.g. fetchNames) y11 and y12 and Graph2 expects to inputs (e.g. feedNames) x21 and x22:
|===
|Graph1:|Graph2:
| y11 = x1 * 2 | y2 = x21 + x22
| y12 = x1 * 3 |
|===
The composed graph would look like this:
[source,Java]
----
Composed = Graph1.andThen( map: y11 -> x21 and y12 -> x22).andThen(Graph2)
----
The https://github.com/spring-cloud/stream-applications/blob/master/functions/common/tensorflow-common/src/main/java/org/springframework/cloud/fn/common/tensorflow/Functions.java#L72[Functions#rename] utility helps to define the input/output mappings as illustrated in the https://github.com/spring-cloud/stream-applications/blob/master/functions/common/tensorflow-common/src/test/java/org/springframework/cloud/fn/common/tensorflow/FunctionCompositionMultipleInputsOutputs.java[FunctionCompositionMultipleInputsOutputs] example:
[source,Java]
----
try (
GraphRunner graph1 = new GraphRunner(Arrays.asList("x1"), Arrays.asList("y11", "y12"))
.withGraphDefinition(tf -> {
Placeholder<Integer> x1 = tf.withName("x1").placeholder(Integer.class);
tf.withName("y11").math.mul(x1, tf.constant(2));
tf.withName("y12").math.mul(x1, tf.constant(3));
});
GraphRunner graph2 = new GraphRunner(Arrays.asList("x21", "x22"), Arrays.asList("y2"))
.withGraphDefinition(tf -> tf.withName("y2").math.add(
tf.withName("x21").placeholder(Integer.class),
tf.withName("x22").placeholder(Integer.class)));
Tensor x = Tensor.create(10);
) {
Map<String, Tensor<?>> result =
graph1
.andThen(
Functions.rename(
"y11", "x21",
"y12", "x22"
))
.andThen(graph2)
.apply(Collections.singletonMap("x", x));
System.out.println("Result is: " + result.get("y2").intValue()); // Result is: 50
}
----
The Functions#rename(String...mappings) takes an even number of string pairs, where every even parameter represents the from and to name to map. Eg. The y11 above is mapped into x21 and y12 is mapped into x22. +
The https://github.com/spring-cloud/stream-applications/blob/master/functions/common/tensorflow-common/src/main/java/org/springframework/cloud/fn/common/tensorflow/AbstractGraphRunner.java#L129[GraphRunner#enableAutoBinding()] and https://github.com/spring-cloud/stream-applications/blob/master/functions/common/tensorflow-common/src/main/java/org/springframework/cloud/fn/common/tensorflow/AbstractGraphRunner.java#L124[GraphRunner#disableAutoBinding()] allow altering the autobinding behavior enforcing mapping even of singleton input/output graphs.
=== Save and Close Obsolete Tensors
The Tensors used as inputs (feeds) and outputs (fetches) by the GraphRunners have to be released (e.g. closed) when not used anymore.
Because every sub-graph in a composite pipeline produces one or more <String, Tensor> pairs we need to track those references and close them.
The https://github.com/spring-cloud/stream-applications/blob/master/functions/common/tensorflow-common/src/main/java/org/springframework/cloud/fn/common/tensorflow/GraphRunnerMemory.java[GraphRunnerMemory] is a handy utility Function implementation that keeps track of all input Tensor parameters passed through. It is https://docs.oracle.com/javase/8/docs/api/java/lang/AutoCloseable.html[AutoClosable] and will release all tracked Tensors when closed.
The GraphRunnerMemory implements the same function signatures as the GraphRunner (e.g. Fun<Map<S,T>, Map<S,T>>) and therefore can participate in composite graph definitions:
[source,Java]
----
try ( memory = new GraphRunnerMemory() ) {
composed-graph =
Graph1..andThen(memory)
.andThen(Graph2).andThen(memory)
…
.andThen(GraphN).andThen(memory)
….
} // releases all Tensors returned by the GraphRunners
----
The https://github.com/spring-cloud/stream-applications/blob/master/functions/common/tensorflow-common/src/test/java/org/springframework/cloud/fn/common/tensorflow/ReleaseTensorParameters.java[ReleaseTensorParameters] example illustrates how to use the GraphRunnerMemory:
[source,Java]
----
try (
Tensor x = Tensor.create(input);
GraphRunnerMemory memory = new GraphRunnerMemory();
) {
Map<String, Tensor<?>> result =
this.graph1.andThen(memory)
.andThen(this.graph2).andThen(memory)
.apply(Collections.singletonMap("x", x));
return result.get("y2").intValue();
}
// At that point all intermediate Tensors used by the GraphRunners are closed.
----
Note: the GraphRunnerMemory has some other very useful applications that we will highlight in the next paragraph.
=== Enrich Graph Inputs
For particular graphs in the composite pipeline, we can add an additional input parameters that were not produced by the upstream graph.
With the help fo the https://github.com/spring-cloud/stream-applications/blob/master/functions/common/tensorflow-common/src/main/java/org/springframework/cloud/fn/common/tensorflow/Functions.java#L42[Functions#enrichWith(name, Tensor)] utility function we can inject the additional parameters in the graph composition.
In the following snippet we enrich the graph2’s input with an additional parameter (newParam):
[source,Java]
----
try (
Tensor x = Tensor.create(input);
Tensor additionalTensor = Tensor.create(colorMap);
) {
Map<String, Tensor<?>> result =
graph1
.andThen(Functions.enrichWIth("newParam", additionalTensor)
.andThen(graph2)
.apply(Collections.singletonMap("x", x));
return result.get("y2").intValue();
}
----
The https://github.com/spring-cloud/stream-applications/blob/master/functions/function/semantic-segmentation-function/src/main/java/org/springframework/cloud/fn/semantic/segmentation/SemanticSegmentation.java#L150[SemanticSegmentation] implementation provides a real example how to enrich with parameters.
=== Enrich Inputs from Saved Tensors
We can combine the enricher approach with the https://github.com/spring-cloud/stream-applications/blob/master/functions/common/tensorflow-common/src/main/java/org/springframework/cloud/fn/common/tensorflow/GraphRunnerMemory.java[GraphRunnerMemory]. This allows us to enrich some downstream Graphs with tensor parameters computed in some of the upstream Graphs. The `Functions#enrichFromMemory(memory, tensorName)` utility function can enrich a graph input parameter by extracting one stored in the memory.
For example let’s construct the following graph compositions:
----
graph1: y1 = x1 * 10 +
graph2: y2 = y1 * 200 +
graph3: y3 = y2 + y1
----
[source,Java]
----
try (
Tensor x = Tensor.create(input);
GraphRunnerMemory memory = new GraphRunnerMemory();
) {
Map<String, Tensor<?>> result =
this.graph1.andThen(memory) // memorizes y1
.andThen(graph2).andThen(memory) // memorizes y2
.andThen(Functions.enrichFromMemory(memory, "y1")) // retrieve graph1’s output y1 and adds it as an input for the next function.
.andThen(Functions.rename(
"y1", "x31", // renames the input y1 into x31
"y2", "x32" // renames the input y2 into x32
))
.andThen(graph3).andThen(memory)
.apply(Collections.singletonMap("x", x));
return result.get("y3").intValue();
}
----
=== Load Frozen Models from Remote Archives
The ProtoBufGraphDefinition extracts a pre-trained (frozen) Tensorflow model form a URI archive into byte array. It supports the `http(s)://`, `file://` and `classpath://` URI schemas. For this it uses the `ModelExtractor` and `CachedModelExtractor` utilities.
Models can be extracted either from raw files or form compressed archives. When extracted from an archive the model file name can optionally be provided as a URI fragment. For example for resource: `http://myarchive.tar.gz#model.pb`
the `myarchive.tar.gz` is traversed to uncompress and extract the model.pb file as a byte array. If the file name is not provided as URI fragment then the first file in the archive with extension .pb is extracted.
In addition, the CachedModelExtractor allows keeping a local copy (cache) of the model (protobuf) files extracted from the URI archive.
|===
|The https://github.com/spring-cloud/stream-applications/tree/master/functions/function/image-recognition-function[image-recognition] and https://github.com/spring-cloud/stream-applications/tree/master/functions/function/semantic-segmentation-function[semantic-segmentation] inference models implementations demonstrate the suggested programming model.
|===
