Note that the lifetime of a function is coupled with the lifetime of its graph or session, i.e. + * the function will failed to be invoked after the graph or session is released, which ever comes + * first. e.g. + * + *
{@code
+ * FunctionGraph function;
+ * try (Graph g = new Graph()) {
+ * Ops tf = Ops.create(g);
+ * Placeholder x = tf.placeholder(TFloat32.DTYPE);
+ * Add y = tf.math.add(x, tf.constant(2.0f));
+ * try (Session s = new Session(s)) {
+ * function = FunctionGraph.builder("myFunction").input("x", x).output("y", y).build(s);
+ * try (Tensor xValue = TFloat32.scalarOf(10.0f);
+ * Tensor yValue = function.call(xValue).expect(TFloat32.DTYPE)) {
+ * assertEquals(12.0f, yValue.data().getFloat());
+ * }
+ * }
+ * }
+ * try (Tensor xValue = TFloat32.scalarOf(10.0f)) {
+ * function.call(xValue); // fails, graph has been closed
+ * }
+ * }
+ *
+ * A function can also invoke a + * tf.function + * defined in a {@link SavedModelBundle}. + * + *
{@code
+ * FunctionGraph myFunction = savedModelBundle.function("myFunctionSignatureName");
+ * Map> outputTensorMap = myFunction.call(inputTensorMap);
+ * }
+ */
+public class FunctionGraph {
+
+ /** The default signature name, when not provided */
+ public static final String DEFAULT_NAME = "serving_default";
+
+ /**
+ * Builds a new function signature.
+ */
+ public static class Builder {
+
+ /**
+ * Register a tensor as an input of the function.
+ *
+ * @param inputName user-friendly name for this input tensor
+ * @param input input tensor
+ * @return this builder
+ */
+ public Builder input(String inputName, Operand input) {
+ signatureBuilder.putInputs(inputName, toTensorInfo(input.asOutput()));
+ return this;
+ }
+
+ /**
+ * Register a tensor as an output of the function.
+ *
+ * @param inputName user-friendly name for this input tensor
+ * @param input input tensor
+ * @return this builder
+ */
+ public Builder output(String outputName, Operand output) {
+ signatureBuilder.putOutputs(outputName, toTensorInfo(output.asOutput()));
+ return this;
+ }
+
+ /**
+ * Provide extensible name information enabling third-party users to mark a signature as
+ * supporting a particular method
+ *
+ * @param methodName method name
+ * @return this builder
+ */
+ public Builder methodName(String methodName) {
+ signatureBuilder.setMethodName(methodName);
+ return this;
+ }
+
+ /**
+ * Creates a function from a graph session.
+ *
+ * The provided session will be used for running or saving this function. + * + * @param signature signature of the function + * @param session a graph session + * @return a function + */ + public FunctionGraph build(Session session) { + return new FunctionGraph(name, signatureBuilder.build(), session); + } + + private static TensorInfo toTensorInfo(Output operand) { + Shape shape = operand.shape(); + TensorShapeProto.Builder tensorShapeBuilder = TensorShapeProto.newBuilder(); + for (int i = 0; i < shape.numDimensions(); ++i) { + tensorShapeBuilder.addDim(Dim.newBuilder().setSize(shape.size(i))); + } + return TensorInfo.newBuilder() + .setDtype(DataType.forNumber(operand.dataType().nativeCode())) + .setTensorShape(tensorShapeBuilder) + .setName(operand.op().name() + ":" + operand.index()) + .build(); + } + + private final String name; + private final SignatureDef.Builder signatureBuilder = SignatureDef.newBuilder(); + + private Builder(String name) { + this.name = name; + } + } + + /** + * Returns a new builder for creating a function + * + *
"serving_default" will be used as the default function signature name.
+ */
+ public static Builder builder() {
+ return new Builder(DEFAULT_NAME);
+ }
+
+ /**
+ * Returns a new builder for creating a function.
+ *
+ * @param name function signature name
+ */
+ public static Builder builder(String name) {
+ return new Builder(name);
+ }
+
+ /**
+ * Return the name of this function
+ */
+ public String name() {
+ return name;
+ }
+
+ /**
+ * Returns the method name of this function (e.g. as exposed by a server)
+ */
+ public String methodName() {
+ return signatureDef.getMethodName();
+ }
+
+ /**
+ * Returns the names of the inputs of this function.
+ */
+ public Set Caller is responsible for closing all Tensors.
+ *
+ * @param tensor input tensor
+ * @return output tensor
+ */
+ public Map Caller is responsible for closing all Tensors.
+ *
+ * @param tensor input tensor
+ * @return output tensor
+ * @throws IllegalArgumentException if there are multiple input or output parameters defined
+ * in the function
+ */
+ public Tensor call(Tensor tensor) throws IllegalArgumentException {
+ if (signatureDef.getInputsCount() != 1) {
+ throw new IllegalArgumentException(
+ String.format("Function [%s] requires multiple inputs", signatureDef.getMethodName()));
+ }
+ String inputNodeName = signatureDef.getInputsMap().values().iterator().next().getName();
+
+ if (signatureDef.getOutputsCount() != 1) {
+ throw new IllegalArgumentException(
+ String.format("Function [%s] has multiple outputs", signatureDef.getMethodName()));
+ }
+ String outputNodeName = signatureDef.getOutputsMap().values().iterator().next().getName();
+
+ return session.runner().feed(inputNodeName, tensor).fetch(outputNodeName).run().get(0);
+ }
+
+ Session session() {
+ return session;
+ }
+
+ SignatureDef signatureDef() {
+ return signatureDef;
+ }
+
+ private final String name;
+ private final Session session;
+ private final SignatureDef signatureDef;
+
+ FunctionGraph(String name, SignatureDef signatureDef, Session session) {
+ this.name = name;
+ this.session = session;
+ this.signatureDef = signatureDef;
+ }
+}
diff --git a/tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/Graph.java b/tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/Graph.java
index 221268191fd..071a28bb6c9 100644
--- a/tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/Graph.java
+++ b/tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/Graph.java
@@ -43,8 +43,17 @@
import org.tensorflow.internal.c_api.TF_Output;
import org.tensorflow.internal.c_api.TF_Status;
import org.tensorflow.internal.c_api.TF_WhileParams;
+import org.tensorflow.ndarray.StdArrays;
import org.tensorflow.op.Op;
+import org.tensorflow.op.Ops;
+import org.tensorflow.op.core.Constant;
+import org.tensorflow.op.core.NoOp;
+import org.tensorflow.op.core.Placeholder;
+import org.tensorflow.op.train.Restore;
+import org.tensorflow.op.train.Save;
import org.tensorflow.proto.framework.GraphDef;
+import org.tensorflow.proto.util.SaverDef;
+import org.tensorflow.types.TString;
/**
@@ -67,6 +76,11 @@ public Graph() {
this.nativeHandle = nativeHandle;
}
+ Graph(TF_Graph nativeHandle, SaverDef saverDef) {
+ this(nativeHandle);
+ this.saverDef = saverDef;
+ }
+
/**
* Release resources associated with the Graph.
*
@@ -287,6 +301,17 @@ public Output[] addGradients(Output y, Output[] x) {
return addGradients(null, new Output[] {y}, x, null);
}
+ public SaverDef saverDef() {
+ if (saverDef == null) {
+ synchronized (this) {
+ if (saverDef == null) {
+ saverDef = addVariableSaver(this);
+ }
+ }
+ }
+ return saverDef;
+ }
+
/**
* Used to instantiate an abstract class which overrides the buildSubgraph method to build a
* conditional or body subgraph for a while loop. After Java 8, this can alternatively be used to
@@ -405,6 +430,7 @@ public Output[] whileLoop(
private final Object nativeHandleLock = new Object();
private TF_Graph nativeHandle;
private int refcount = 0;
+ private SaverDef saverDef;
private final List Caller is responsible for closing all returned Tensors.
+ *
+ * @param arguments list of input tensors, mapped by their signature name
+ * @return list of output tensors, mapped by the signature name
+ * @throws IllegalArgumentException if no function can be selected by default
*/
- public TfFunction function(String functionSignatureName) {
- return new TfFunction(functionSignatureName, this.getSignatureToNodeName(), this.session);
+ public Map Invoked from the native load method. Takes ownership of the handles.
*/
private static SavedModelBundle fromHandle(
- TF_Graph graphHandle, TF_Session sessionHandle, MetaGraphDef metaGraphDef) {
- Graph graph = new Graph(graphHandle);
- Session session = new Session(graph, sessionHandle);
- return new SavedModelBundle(graph, session, metaGraphDef);
+ final TF_Graph graphHandle, final TF_Session sessionHandle, MetaGraphDef metaGraphDef) {
+
+ final Graph graph = new Graph(graphHandle, metaGraphDef.getSaverDef());
+ final Session session = new Session(graph, sessionHandle);
+
+ // For each signature definition, create a distinct function based on the main graph/session
+ final MapSavedModelBundle with the configured options. */
public SavedModelBundle load() {
return SavedModelBundle.load(exportDir, tags, configProto, runOptions);
@@ -98,99 +113,99 @@ private Loader(String exportDir) {
private RunOptions runOptions = null;
}
- /**
- * SignatureToNodeName finds the node names in the {@link Graph} corresponding to the
- * input / output parameters of a tf.function
- */
- public static final class SignatureToNodeName {
-
- public SignatureToNodeName(SavedModelBundle savedModelBundle) {
- loadSignatures(savedModelBundle);
- }
+ /** Options for exporting a SavedModel. */
+ public static final class Exporter {
/**
- * Given a tf.function signature name, find the node names corresponding
- * to the input arguments
+ * Sets the set of tags that identify the specific graph in the saved model to save.
+ *
+ * Note that only one graph per model can be saved right now using this API.
*
- * @param functionSignatureName tf.function signature name
- * @return a map from input arguments to node names in the {@link Graph}
+ * @param tags the tags identifying the specific MetaGraphDef to save.
+ * @return this object
*/
- public MapExporter object for setting configuration options before actually
+ * saving the model.
+ *
+ * @param exportDir the directory path containing a saved model.
+ */
+ public static Exporter exporter(String exportDir) {
+ return new Exporter(exportDir);
+ }
+
/**
* Returns the MetaGraphDef
@@ -248,31 +275,54 @@ public Session session() {
}
/**
- * Returns the {@link SignatureToNodeName} translator for the model.
+ * Return a {@link FunctionGraph} corresponding to the function signature.
+ *
+ * {@code
+ * FunctionGraph myFunction = savedModelBundle.function("myFunctionSignatureName");
+ * Map
*
- * @return SignatureToNodeName translator
+ * @param functionSignatureName name of the {@code SignatureDef} in the saved model.
+ * @return TfFunction object that can be used to make calls to the tf.function
+ * @throws IllegalArgumentException if {@code functionSignatureName} is not found in this
+ * saved model.
*/
- public SignatureToNodeName getSignatureToNodeName() {
- if (this.sigToNodeName == null) {
- // no need to lock, ok to create multiple instances
- this.sigToNodeName = new SignatureToNodeName(this);
+ public FunctionGraph function(String functionSignatureName) {
+ FunctionGraph function = functions.get(functionSignatureName);
+ if (function == null) {
+ throw new IllegalArgumentException(
+ String.format("Function with signature [%s] not found", functionSignatureName));
}
- return this.sigToNodeName;
+ return function;
}
/**
- * Return a {@link TfFunction} corresponding to the function signature.
+ * Invokes the default function directly from this model.
*
- * {@code
- * TfFunction myFunction = savedModelBundle.function("myFunctionSignatureName");
- * Map
+ * The default function selection is done based on the first of the following conditions that
+ * is true:
+ *
+ *
*
- * @param functionSignatureName name of the {@code SignatureDef} in the saved model.
- * @return TfFunction object that can be used to make calls to the tf.function
+ * {@code
- * TfFunction myFunction = savedModelBundle.function("myFunctionSignatureName");
- * Map
- *
- */
-public class TfFunction {
-
- public TfFunction(
- String functionSignatureName,
- SavedModelBundle.SignatureToNodeName nameToNode, Session session) {
- this.nameToNode = nameToNode;
- this.session = session;
- this.functionSignatureName = functionSignatureName;
- }
-
- /**
- * Invokes a tf.function.
- * Caller is responsible for closing all Tensors.
- *
- * @param arguments map of input tensors
- * @return map of output tensors
- */
- public Map