Execute graph initializers in a single call

Author: karllessard

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

@@ -183,13 +183,20 @@ public synchronized void addInitializer(Op initializer) {
   }
 
   /**
-   * Returns an op which initializers all the variables.
-   * @return The initializer operation.
+   * Create an op which executes all registered initializers in the graph.
+   * @return the initializer op
    */
   public NoOp variablesInitializer() {
     return variablesInitializer(DEFAULT_INIT_NAME);
   }
 
+  /**
+   * Returns an op with the given {@code name} which executes all registered initializers in the
+   * graph.
+   *
+   * @param name name to give to the initializer op
+   * @return the initializer op
+   */
   public NoOp variablesInitializer(String name) {
     Scope scope = new Scope(this);
     scope = scope.withName(name).withControlDependencies(initializers);

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

@@ -36,6 +36,7 @@
 import org.tensorflow.internal.c_api.TF_SessionOptions;
 import org.tensorflow.internal.c_api.TF_Status;
 import org.tensorflow.internal.c_api.TF_Tensor;
+import org.tensorflow.op.Op;
 
 /**
  * Driver for {@link Graph} execution.
@@ -432,6 +433,42 @@ public Runner runner() {
     return new Runner();
   }
 
+  /**
+   * Run all graph initializers.
+   *
+   * <p>Initializers must be executed once before running the graph using a session
+   * {@link Runner}.</p>
+   *
+   * <p>This method can be used only if the initializers were added programatically to the graph
+   * via {@link Graph#addInitializer(Op)}. Otherwise, {@link #runInitializers(String)} must be used
+   * to provide the name of the initializer operation to execute.
+   */
+  public void runInitializers() {
+    runner().addTarget(graph.variablesInitializer().op()).run();
+  }
+
+  /**
+   * Run all graph initializers grouped under the {@code initializerOpName} operation.
+   *
+   * <p>Initializers must be executed once before running the graph using a session
+   * {@link Runner}.</p>
+   *
+   * <p>The {@code initializerOpName} is the name of a single operation already registered to the
+   * graph that executes all graph initializers at once. If the initializers were created using
+   * {@link Graph#variablesInitializer(String)}, the names in both methods must match. Otherwise,
+   * {@link Graph#DEFAULT_INIT_NAME} should be used.</p>
+   *
+   * @param initializerOpName name of the initializer operation.
+   */
+  public void runInitializers(String initializerOpName) {
+    Operation operation = graph.operation(initializerOpName);
+    if (operation == null) {
+      throw new IllegalArgumentException("Initializer operation named '"
+          + initializerOpName + "' cannot be found in the graph");
+    }
+    runner().addTarget(operation).run();
+  }
+
   /**
    * Output tensors and metadata obtained when executing a session.
    *

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

@@ -24,7 +24,9 @@
 import org.junit.runners.JUnit4;
 import org.tensorflow.op.Ops;
 import org.tensorflow.op.core.Split;
+import org.tensorflow.op.core.Variable;
 import org.tensorflow.op.linalg.MatMul;
+import org.tensorflow.op.math.Add;
 import org.tensorflow.tools.Shape;
 import org.tensorflow.tools.ndarray.NdArrays;
 import org.tensorflow.tools.ndarray.StdArrays;
@@ -161,6 +163,64 @@ public void createWithConfigProto() {
         Session s = new Session(g, singleThreadConfigProto())) {}
   }
 
+  @Test
+  public void runInitializers() {
+    try (Graph g = new Graph()) {
+      Ops tf = Ops.create(g);
+
+      Variable<TInt32> var1 = tf.variable(Shape.scalar(), TInt32.DTYPE);
+      Variable<TInt32> var2 = tf.variable(Shape.scalar(), TInt32.DTYPE);
+      Add<TInt32> add = tf.math.add(var1, var2);
+
+      g.addInitializer(tf.assign(var1, tf.constant(10)));
+      g.addInitializer(tf.assign(var2, tf.constant(20)));
+
+      try (Session s = new Session(g)) {
+        s.runInitializers();
+
+        try (Tensor<TInt32> t = s.runner().fetch(add).run().get(0).expect(TInt32.DTYPE)) {
+          assertEquals(30, t.data().getInt());
+        }
+      }
+    }
+  }
+
+  @Test
+  public void runInitializersByName() {
+    try (Graph g = new Graph()) {
+      Ops tf = Ops.create(g);
+
+      Variable<TInt32> var1 = tf.variable(Shape.scalar(), TInt32.DTYPE);
+      Variable<TInt32> var2 = tf.variable(Shape.scalar(), TInt32.DTYPE);
+      Add<TInt32> add = tf.math.add(var1, var2);
+
+      g.addInitializer(tf.assign(var1, tf.constant(10)));
+      g.addInitializer(tf.assign(var2, tf.constant(20)));
+      g.variablesInitializer("init_test");
+      try {
+        g.variablesInitializer("init_test");
+        fail();
+      } catch (IllegalArgumentException e) {
+        // as expected, cannot register initializers twice
+      }
+
+      try (Session s = new Session(g)) {
+        s.runInitializers("init_test");
+
+        try (Tensor<TInt32> t = s.runner().fetch(add).run().get(0).expect(TInt32.DTYPE)) {
+          assertEquals(30, t.data().getInt());
+        }
+
+        try {
+          s.runInitializers("wrong_name");
+          fail();
+        } catch (IllegalArgumentException e) {
+          // as expected
+        }
+      }
+    }
+  }
+
   private static byte[] fullTraceRunOptions() {
     // Ideally this would use the generated Java sources for protocol buffers
     // and end up with something like the snippet below. However, generating