Fix filename tensor name in SaverDef

Author: karllessard

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

@@ -461,11 +461,12 @@ synchronized SaverDef saverDef() {
         // regenerate SaverDef without mutating. The names mirror
         // the python implementation for compatibility.
         // https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/training/saver.py
-        saverDef = SaverDef.newBuilder()
-                           .setFilenameTensorName("save/filename")
-                           .setSaveTensorName("save/control_dependency")
-                           .setRestoreOpName("save/restore_all")
-                           .build();
+        saverDef =
+            SaverDef.newBuilder()
+                .setFilenameTensorName("save/filename:0")
+                .setSaveTensorName("save/control_dependency")
+                .setRestoreOpName("save/restore_all")
+                .build();
       }
     }
     return saverDef;
@@ -812,36 +813,35 @@ private static SaverDef addVariableSaver(Graph graph) {
       }
     }
 
-    // FIXME Need an easier way to initialize an NdArray from a list
-    String[] tmp = new String[varNames.size()];
-    Constant<TString> varNamesTensor = tf.constant(StdArrays.ndCopyOf(varNames.toArray(tmp)));
-    Operand<TString> varSlices = tf.zerosLike(varNamesTensor);
-
-    Placeholder<TString> saveFilename = tf.withName("filename").placeholder(TString.class);
-    Save saveVariables = tf.train.save(
-        saveFilename,
-        varNamesTensor,
-        varSlices,
-        varOutputs
-    );
-    Identity<TString> id = tf.withControlDependencies(Arrays.asList(saveFilename,saveVariables))
-            .withName("control_dependency").identity(saveFilename);
-    Restore restoreVariables = tf.train.restore(
-        saveFilename,
-        varNamesTensor,
-        varSlices,
-        varTypes
-    );
-    List<Op> restoreOps = new ArrayList<>(varOutputs.size());
-    for (int i = 0; i < varOutputs.size(); ++i) {
-      restoreOps.add(tf.assign(varOutputs.get(i), (Operand) restoreVariables.tensors().get(i)));
+    Placeholder<TString> filename = tf.withName("filename").placeholder(TString.class);
+    Identity<TString> save = null;
+    NoOp restore = null;
+
+    if (varNames.isEmpty()) {
+      save = tf.withName("empty_save").identity(filename);
+      restore = tf.withName("restore_all").noOp();
+    } else {
+      String[] tmp = new String[varNames.size()];
+      Constant<TString> varNamesTensor = tf.constant(StdArrays.ndCopyOf(varNames.toArray(tmp)));
+      Operand<TString> varSlices = tf.zerosLike(varNamesTensor);
+      Save saveVars = tf.train.save(filename, varNamesTensor, varSlices, varOutputs);
+      List<Op> saveDeps = Arrays.asList(filename, saveVars);
+      Restore restoreVars = tf.train.restore(filename, varNamesTensor, varSlices, varTypes);
+      List<Op> restoreDeps = new ArrayList<>(varOutputs.size());
+      for (int i = 0; i < varOutputs.size(); ++i) {
+        restoreDeps.add(tf.assign(varOutputs.get(i), (Operand) restoreVars.tensors().get(i)));
+      }
+      save = tf.withControlDependencies(saveDeps).withName("control_dependency").identity(filename);
+      restore = tf.withControlDependencies(restoreDeps).withName("restore_all").noOp();
     }
-    NoOp restoreAll = tf.withControlDependencies(restoreOps).withName("restore_all").noOp();
 
+    // 'Filename' must be the name of a tensor (i.e. with output index)
+    // 'Save' must be an operation name, even if the field name is confusing (see SaverDef doc)
+    // 'Restore' must be an operation name
     return SaverDef.newBuilder()
-        .setFilenameTensorName(saveFilename.op().name())
-        .setSaveTensorName(id.op().name())
-        .setRestoreOpName(restoreAll.op().name())
+        .setFilenameTensorName(filename.output().name())
+        .setSaveTensorName(save.op().name())
+        .setRestoreOpName(restore.op().name())
         .build();
   }
 

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

@@ -38,6 +38,11 @@ public int index() {
     return index;
   }
 
+  /** Returns the full name of this Output (a.k.a. tensor name) */
+  public String name() {
+    return op().name() + ":" + index;
+  }
+
   /** Returns the DataType of the tensor referred to by this Output. */
   @SuppressWarnings("unchecked")
   public DataType dataType() {
@@ -48,7 +53,7 @@ public DataType dataType() {
   @SuppressWarnings("unchecked")
   @Override
   public Class<T> type() {
-    return (Class<T>)TensorTypeRegistry.find(dataType()).type();
+    return (Class<T>) TensorTypeRegistry.find(dataType()).type();
   }
 
   /**
@@ -63,7 +68,10 @@ public Class<T> type() {
   public <U extends TType> Output<U> expect(Class<U> type) {
     if (type != type()) {
       throw new IllegalArgumentException(
-          "Cannot cast from output of " + this.type().getSimpleName() + " to output of " + type.getSimpleName());
+          "Cannot cast from output of "
+              + this.type().getSimpleName()
+              + " to output of "
+              + type.getSimpleName());
     }
     return ((Output<U>) this);
   }
@@ -80,17 +88,16 @@ public <U extends TType> Output<U> expect(Class<U> type) {
    *
    * @return tensor
    * @throws IllegalStateException if this output results from a graph
-   * @throws ClassCastException if the type of the tensor and this output are unexpectedly incompatible
+   * @throws ClassCastException if the type of the tensor and this output are unexpectedly
+   *     incompatible
    * @see EagerSession
    */
   @SuppressWarnings("unchecked")
   public T asTensor() {
-    return (T)operation.tensor(index);
+    return (T) operation.tensor(index);
   }
 
-  /**
-   * Returns the (possibly partially known) shape of the tensor referred to by this output.
-   */
+  /** Returns the (possibly partially known) shape of the tensor referred to by this output. */
   @Override
   public Shape shape() {
     return operation.shape(index);

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

@@ -1,18 +1,18 @@
-/*
- * Copyright 2020 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.
- */
+/* Copyright 2020-2021 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 java.util.HashMap;
@@ -171,7 +171,7 @@ public Set<String> outputNames() {
 
   @Override
   public String toString() {
-    StringBuilder strBuilder = new StringBuilder("Signature for \"" + key +"\":\n");
+    StringBuilder strBuilder = new StringBuilder("Signature for \"" + key + "\":\n");
     String methodName = methodName();
     if (methodName != null && !methodName.isEmpty()) {
       strBuilder.append("\tMethod: \"").append(methodName).append("\"\n");

tensorflow-core/tensorflow-core-api/src/test/java/org/tensorflow/SavedModelBundleTest.java

@@ -44,6 +44,7 @@
 import org.tensorflow.proto.framework.RunOptions;
 import org.tensorflow.proto.framework.SignatureDef;
 import org.tensorflow.proto.framework.TensorInfo;
+import org.tensorflow.proto.util.SaverDef;
 import org.tensorflow.types.TFloat32;
 
 /** Unit tests for {@link org.tensorflow.SavedModelBundle}. */
@@ -123,7 +124,13 @@ public void exportFunctionWithVariables() throws IOException {
     try (SavedModelBundle savedModel =
         SavedModelBundle.load(testFolder.toString(), SavedModelBundle.DEFAULT_TAG)) {
       assertNotNull(savedModel.metaGraphDef());
-      assertNotNull(savedModel.metaGraphDef().getSaverDef());
+
+      SaverDef saverDef = savedModel.metaGraphDef().getSaverDef();
+      assertNotNull(saverDef);
+      assertEquals("save/filename:0", saverDef.getFilenameTensorName());
+      assertEquals("save/control_dependency", saverDef.getSaveTensorName());
+      assertEquals("save/restore_all", saverDef.getRestoreOpName());
+
       assertEquals(1, savedModel.metaGraphDef().getSignatureDefCount());
       assertEquals(Signature.DEFAULT_KEY,
           savedModel.metaGraphDef().getSignatureDefMap().keySet().iterator().next());