Save models as functions

[karllessard]

Ok I've modified a little this solution compared to the previous one, now FunctionGraph are auto-closeable and own the session/graph only if the function has been created before them. See examples provided in the FunctionGraph documentation for more details.

There are still a few things I need to look at before making this PR more official but I would like to get some feedbacks before continuing in that direction.

In summary:

• FunctionGraph is the new (disposable) name for TfFunction as submitted previously by @Shajan
• Users can create their graph in a "function style" (while we still don't support real function graphs but this API should be ready to accommodate it when we will come to it).
• Users can still define a graph and then convert it to a function (i.e. adding a signature to it). In this case, the function does not own the graph and releasing it will have no effect.
• Saving a model requires that one or more functions to be provided. The graph to be saved and the session to use to save the variables will be extracted from these functions (right now, since we can only export session-centric models, all functions saved should share the same session)

I'll provide more and better examples later, thanks!

[googlebot]

All (the pull request submitter and all commit authors) CLAs are signed, but one or more commits were authored or co-authored by someone other than the pull request submitter.

We need to confirm that all authors are ok with their commits being contributed to this project. Please have them confirm that by leaving a comment that contains only @googlebot I consent. in this pull request.

Note to project maintainer: There may be cases where the author cannot leave a comment, or the comment is not properly detected as consent. In those cases, you can manually confirm consent of the commit author(s), and set the cla label to yes (if enabled on your project).

ℹ️ Googlers: Go here for more info.

[Shajan]

Suggestion: Rename GRAPH to GRAPH_AND_SESSION

[Shajan]

The change looks good.

[Craigacp]

A few comments on names and other small things.

[Craigacp]

Maybe FunctionSession as it's a specialised session? That way it conceptually lives next to EagerSession and Session, rather than next to Graph. I feel like this is much closer to a session than a graph.

[karllessard]

Right now, it is true that all functions loaded from a saved model share the same graph and are just used to execute it with a given signature (therefore they are acting more as a session). Nonetheless, I still have a preference to preserve the FunctionGraph naming.

If we ignore the saving part (which is limited due to the actual state of the C API), conceptually these functions also allow you to build your graphs, replacing the need of allocating explicitly a Graph instance, e.g.

public class MyModel {

  // Function graph #1
  private static Signature addTwo(Ops tf) {
    Placeholder<TFloat32> input = tf.placeholder(TFloat32.DTYPE);
    Add<TFloat32> output = tf.math.add(input, tf.constant(2.0f));
    return Signature.builder("addTwo").input("x", input).output("y", output).build();
  }

  // Function graph #2
  private static Signature minusTwo(Ops tf) {
    Placeholder<TFloat32> input = tf.placeholder(TFloat32.DTYPE);
    Add<TFloat32> output = tf.math.sub(input, tf.constant(2.0f));
    return Signature.builder("subTwo").input("x", input).output("y", output).build();
  }

  public static void main(String args[]) {
    try (FunctionGraph f1 = FunctionGraph.create(MyModel::addTwo);
         FunctionGraph f2 = FunctionGraph.create(MyModel::minusTwo);
         Tensor<TFloat32> x = TFloat32.scalarOf(3.0f)) {
        assertEquals(3.0f, f2.call(f1.call(x)).expect(TFloat32.DTYPE).data().getFloat());
    }
  }
}

Each function has its own graph and therefore if very coupled with that concept (where EagerSession has no graph at all). They will also appear as separate graphs when exporting them into a saved model the same way Python does it (in fact, Python saves each function as one or more "objects", which are then linked to their distinct graph in the function library).

[karllessard]

Another idea would be to call this class a ConcreteFunction. Concrete functions and functions are two very similar but distinct concepts in TF Python. The former is a typed realization of a function graph while the latter is its polymorphic version, only acting as a facade (e.g. a function is eventually composed of multiple concrete functions, one for each type of operands the function has been called with).

In our scenario here, the function graphs are strongly typed and refer to a single graph, therefore more behave like a concrete function in the Python paradigm and it would make sense to name them after it.

We can probably support polymorphic functions too in the future, where we pass the type of the input tensors in parameter to the function builder. So we will still need to find a proper name for this concept, which will probably result in class encapsulating one or more ConcreteFunction. Maybe simply PolymorphicFunction?

[Craigacp]

My preference for a non-graph name is that Graphs aren't executable, and are missing necessary state because the Variables live in the Session. So FunctionGraph is more like a Session because it contains all the necessary bits to execute. However I'm also fine with the ConcreteFunction name as that's just a new concept that people have to learn.

[Craigacp]

saverDef needs to be declared volatile for this to work - http://www.cs.umd.edu/~pugh/java/memoryModel/DoubleCheckedLocking.html

[Craigacp]

Or we make the whole method synchronized, which also seems fine as it's unlikely to be heavily contended.

[Craigacp]

Should we add a helper that takes an exportDir, a FunctionGraph and then performs all the necessary operations? Seems like there are a bunch of hoops to jump through if you just want to save a single FunctionGraph.

[karllessard]

I guess we could, what about making that endpoint available at the function level? Something like function.save(exportDir)

[Craigacp]

I like that.

[Craigacp]

Wrong import.

[karllessard]

The change looks good.

Thanks @Shajan , I'll clean this up and merge it to our shared branch. It would be great if you can then validate that this merged solution still successfully cover your initial use case. If so, we will then make an official PR to the master branch out of it.

[karllessard]

@googlebot I consent.

[yzhuang]

Minor feature request: adding a toString() function to this ConcreteFunction class that prints a readable text representation of the signature would be really useful.