Add ability to load ops when loading a Graph

nikhilk · Jan 10, 2018 · 5b2a165 · 5b2a165
1 parent 95bc4b7
commit 5b2a165
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 4 deletions.
diff --git a/samples/graphs/matrix/main.js b/samples/graphs/matrix/main.js
@@ -1,7 +1,6 @@
 const tf = require('tensorflow');
 
-let session = tf.Session.fromGraphDef('./graph.proto', true);
-session.graph.loadOperations({ result: 'computation/result' });
+let session = tf.Session.fromGraphDef('./graph.proto', { result: 'computation/result' });
 
 let outputs = session.run(null, ['result'], null);
 console.log(outputs.result.toValue());

diff --git a/src/graph.js b/src/graph.js
@@ -61,9 +61,15 @@ class Graph extends api.Reference {
     return unresolvedOps;
   }
 
-  static fromGraphDef(graphDefPath) {
+  static fromGraphDef(graphDefPath, operations) {
     let protobuf = fs.readFileSync(graphDefPath);
-    return loadGraph(protobuf);
+    let graph = loadGraph(protobuf);
+
+    if (operations) {
+      graph.loadOperations(operations);
+    }
+
+    return graph;
   }
 
   // TODO: Implement loading a Graph from an in-memory JSON object representatio of a GraphDef.