Adds a closeable session result

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/ConcreteFunction.java

@@ -1,4 +1,4 @@
-/* Copyright 2020-2021 The TensorFlow Authors. All Rights Reserved.
+/* Copyright 2020-2022 The TensorFlow Authors. All Rights Reserved.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -295,8 +295,8 @@ public Operand<?> call(Scope scope, Operand<?> argument) {
   }
 
   @Override
-  public Map<String, Tensor> call(Map<String, Tensor> arguments) {
-    // FIXME need to manage input/output operand lifetimes
+  public Result call(Map<String, Tensor> arguments) {
+    // FIXME need to manage input operand lifetimes
     Ops tf = Ops.create();
     Map<String, Operand<?>> inputs = new LinkedHashMap<>(arguments.size());
 
@@ -305,11 +305,11 @@ public Map<String, Tensor> call(Map<String, Tensor> arguments) {
       inputs.put(inputName, tf.constantOf((TType) argument));
     }
     Map<String, Operand<?>> outputs = tf.call(this, inputs);
-    Map<String, Tensor> tensorOutputs = new LinkedHashMap<>(outputs.size());
+    LinkedHashMap<String, Tensor> tensorOutputs = new LinkedHashMap<>(outputs.size());
     for (String outputName : outputs.keySet()) {
       tensorOutputs.put(outputName, outputs.get(outputName).asTensor());
     }
-    return tensorOutputs;
+    return new Result(tensorOutputs);
   }
 
   /**

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/Result.java

@@ -0,0 +1,198 @@
+/*
+Copyright (c) 2022, Oracle and/or its affiliates. All rights reserved.
+Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+=======================================================================
+*/
+package org.tensorflow;
+
+import org.tensorflow.exceptions.TensorFlowException;
+import org.tensorflow.proto.framework.RunMetadata;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Iterator;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+
+/**
+ * An {@link AutoCloseable} wrapper around a {@link Map} containing {@link Tensor}s.
+ *
+ * <p>When this is closed it closes all the {@link Tensor}s inside it. If you maintain a
+ * reference to a value after this object has been closed it will throw an {@link
+ * IllegalStateException} upon access.
+ *
+ * <p>This class is not thread-safe with respect to the close operation. Multiple closers
+ * or one thread closing a tensor while another is reading may throw exceptions.
+ *
+ * <p>Note this class is used to manage the lifetimes of tensors produced by the
+ * TensorFlow runtime, from sessions and function calls. It is not used as an argument
+ * to {@code session.run} or function calls as users are in control of the creation
+ * of input tensors.
+ */
+public final class Result implements AutoCloseable, Iterable<Map.Entry<String, Tensor>> {
+    @Override
+    public void close() {
+        if (!closed) {
+            for (Tensor t : list) {
+                try {
+                    t.close();
+                } catch (TensorFlowException e) {
+                    logger.log(Level.WARNING, "Exception raised when closing tensor inside result.", e);
+                }
+            }
+            closed = true;
+        } else {
+            logger.warning("Closing an already closed Result");
+        }
+    }
+
+    @Override
+    public Iterator<Map.Entry<String, Tensor>> iterator() {
+        if (!closed) {
+            return map.entrySet().iterator();
+        } else {
+            throw new IllegalStateException("Result is closed");
+        }
+    }
+
+    /**
+     * Returns the number of outputs in this Result.
+     *
+     * @return The number of outputs.
+     */
+    public int size() {
+        return map.size();
+    }
+
+    /**
+     * Gets the set containing all the tensor names.
+     * @return The tensor names set.
+     */
+    public Set<String> keySet() {
+        return Collections.unmodifiableSet(map.keySet());
+    }
+
+    /**
+     * Does this result object have a tensor for the supplied key?
+     * @param key The key to check.
+     * @return True if this result object has a tensor for this key.
+     */
+    public boolean containsKey(String key) {
+        return map.containsKey(key);
+    }
+
+    /**
+     * Gets the value from the container at the specified index.
+     *
+     * <p>Throws {@link IllegalStateException} if the container has been closed, and {@link
+     * IndexOutOfBoundsException} if the index is invalid.
+     *
+     * @param index The index to lookup.
+     * @return The value at the index.
+     */
+    public Tensor get(int index) {
+        if (!closed) {
+            return list.get(index);
+        } else {
+            throw new IllegalStateException("Result is closed");
+        }
+    }
+
+    /**
+     * Gets the value from the container assuming it's not been closed.
+     *
+     * <p>Throws {@link IllegalStateException} if the container has been closed.
+     *
+     * @param key The key to lookup.
+     * @return Optional.of the value if it exists.
+     */
+    public Optional<Tensor> get(String key) {
+        if (!closed) {
+            return Optional.ofNullable(map.get(key));
+        } else {
+            throw new IllegalStateException("Result is closed");
+        }
+    }
+
+    /**
+     * Metadata about the run.
+     *
+     * <p>A <a
+     * href="https://www.tensorflow.org/code/tensorflow/core/protobuf/config.proto">RunMetadata
+     * protocol buffer</a>.
+     */
+    public Optional<RunMetadata> getMetadata() {
+        return Optional.ofNullable(metadata);
+    }
+
+    /**
+     * Creates a Result from the names and values produced by {@link Session.Runner#run()}.
+     *
+     * @param names    The output names.
+     * @param values   The output values.
+     * @param metadata The run metadata, may be null.
+     */
+    Result(List<String> names, List<Tensor> values, RunMetadata metadata) {
+        this.map = new LinkedHashMap<>();
+        this.list = new ArrayList<>(values);
+
+        if (names.size() != values.size()) {
+            throw new IllegalArgumentException(
+                    "Expected same number of names and values, found names.length = "
+                            + names.size()
+                            + ", values.length = "
+                            + values.size());
+        }
+
+        for (int i = 0; i < names.size(); i++) {
+            Tensor old = this.map.put(names.get(i), values.get(i));
+            if (old != null) {
+                throw new IllegalArgumentException("Name collision in the result set, two outputs are named '" + names.get(i) + "'");
+            }
+        }
+        this.metadata = metadata;
+        this.closed = false;
+    }
+
+    /**
+     * Creates a Result from the names and values.
+     *
+     * @param outputs    The run outputs.
+     */
+    Result(LinkedHashMap<String,Tensor> outputs) {
+        this.map = outputs;
+        this.list = new ArrayList<>(outputs.size());
+        for (Map.Entry<String, Tensor> e : outputs.entrySet()) {
+            list.add(e.getValue());
+        }
+        this.metadata = null;
+        this.closed = false;
+    }
+
+    private final Map<String, Tensor> map;
+
+    private final List<Tensor> list;
+
+    private final RunMetadata metadata;
+
+    private boolean closed;
+
+    private static final Logger logger = Logger.getLogger(Result.class.getName());
+}

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/SavedModelBundle.java

@@ -468,7 +468,7 @@ public List<SessionFunction> functions() {
    * @return list of output tensors, mapped by the signature name
    * @throws IllegalArgumentException if no function can be selected by default
    */
-  public Map<String, Tensor> call(Map<String, Tensor> arguments) {
+  public Result call(Map<String, Tensor> arguments) {
     SessionFunction function = null;
     if (functions.size() == 1) {
       function = functions.values().iterator().next();

tensorflow-core/tensorflow-core-api/src/main/java/org/tensorflow/Session.java

@@ -1,4 +1,4 @@
-/* Copyright 2019-2021 The TensorFlow Authors. All Rights Reserved.
+/* Copyright 2019-2022 The TensorFlow Authors. All Rights Reserved.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -21,12 +21,14 @@
 import static org.tensorflow.internal.c_api.global.tensorflow.TF_SetConfig;
 
 import com.google.protobuf.InvalidProtocolBufferException;
+
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+
 import org.bytedeco.javacpp.BytePointer;
 import org.bytedeco.javacpp.Pointer;
 import org.bytedeco.javacpp.PointerPointer;
@@ -306,7 +308,9 @@ public Runner feed(Operand<?> operand, Tensor t) {
      * @throws IllegalArgumentException if no output exists with the provided name
      */
     public Runner fetch(String operation) {
-      return fetch(graph.outputOrThrow(operation));
+      Runner r = fetch(graph.outputOrThrow(operation),false);
+      outputNames.add(operation);
+      return r;
     }
 
     /**
@@ -336,6 +340,20 @@ public Runner fetch(String operation, int index) {
      * @return this session runner
      */
     public Runner fetch(Output<?> output) {
+      return fetch(output, true);
+    }
+
+    /**
+     * Makes {@link #run()} return the Tensor referred to by {@code output}.
+     *
+     * <p>If {@code output} is a resource variable, will fetch the value.
+     *
+     * @param output the node to fetch the tensor from
+     * @param recordName Records the output name. If false the output name must be recorded by the
+     *                   calling method as otherwise the result object will throw on construction.
+     * @return this session runner
+     */
+    private Runner fetch(Output<?> output, boolean recordName) {
       if (output.env() != graph) {
         throw new IllegalStateException(
             "Can't fetch output "
@@ -378,6 +396,9 @@ public Runner fetch(Output<?> output) {
       } else {
         outputs.add(output);
       }
+      if (recordName) {
+        outputNames.add(output.name());
+      }
       return this;
     }
 
@@ -490,13 +511,13 @@ private void doInit() {
      *
      * @return list of resulting tensors fetched by this session runner
      */
-    public List<Tensor> run() {
+    public Result run() {
       doInit();
       return runNoInit();
     }
 
-    List<Tensor> runNoInit() {
-      return runHelper(false).outputs;
+    Result runNoInit() {
+      return runHelper(false);
     }
 
     /**
@@ -509,12 +530,12 @@ List<Tensor> runNoInit() {
      *
      * @return list of resulting tensors fetched by this session runner, with execution metadata
      */
-    public Run runAndFetchMetadata() {
+    public Result runAndFetchMetadata() {
       doInit();
       return runHelper(true);
     }
 
-    private Run runHelper(boolean wantMetadata) {
+    private Result runHelper(boolean wantMetadata) {
       TF_Tensor[] inputTensorHandles = new TF_Tensor[inputTensors.size()];
       TF_Operation[] inputOpHandles = new TF_Operation[inputs.size()];
       int[] inputOpIndices = new int[inputs.size()];
@@ -569,10 +590,7 @@ private Run runHelper(boolean wantMetadata) {
       } finally {
         runRef.close();
       }
-      Run ret = new Run();
-      ret.outputs = outputs;
-      ret.metadata = metadata;
-      return ret;
+      return new Result(outputNames,outputs,metadata);
     }
 
     private class Reference implements AutoCloseable {
@@ -602,6 +620,7 @@ public void close() {
     private final ArrayList<Output<?>> inputs = new ArrayList<>();
     private final ArrayList<Tensor> inputTensors = new ArrayList<>();
     private final ArrayList<Output<?>> outputs = new ArrayList<>();
+    private final ArrayList<String> outputNames = new ArrayList<>();
     private final ArrayList<GraphOperation> targets = new ArrayList<>();
     private RunOptions runOptions = null;
   }
@@ -648,8 +667,9 @@ public SessionFunction function(Signature signature) {
    *
    * @param signature the signature of the function
    * @param arguments the arguments to call with.
+   * @return The results of the function call.
    */
-  public Map<String, Tensor> run(Signature signature, Map<String, Tensor> arguments) {
+  public Result run(Signature signature, Map<String, Tensor> arguments) {
     return function(signature).call(arguments);
   }
 
@@ -698,26 +718,6 @@ public void restore(String prefix) {
     setInitialized();
   }
 
-  /**
-   * Output tensors and metadata obtained when executing a session.
-   *
-   * <p>See {@link Runner#runAndFetchMetadata()}
-   */
-  public static final class Run {
-
-    /** Tensors from requested fetches. */
-    public List<Tensor> outputs;
-
-    /**
-     * Metadata about the run.
-     *
-     * <p>A <a
-     * href="https://www.tensorflow.org/code/tensorflow/core/protobuf/config.proto">RunMetadata
-     * protocol buffer</a>.
-     */
-    public RunMetadata metadata;
-  }
-
   Graph graph() {
     return graph;
   }