Refactor graph loading in tensorflow formatter.
@@ -23,13 +23,22 @@ class TFConstantGraphFormatter(Formatter):
return get_path(path).endswith(".pb")
@classmethod
- def load(cls, path, graph_def=None):
+ def load(cls, path, graph=None, graph_def=None, as_text=False):
import tensorflow as tf
- from tensorflow.python.platform import gfile
+ if not graph:
+ graph = tf.Graph()
+
+ with graph.as_default():
if not graph_def:
graph_def = tf.GraphDef()
+ if as_text:
+ from google.protobuf import text_format
+ with open(get_path(path), "r") as f:
+ text_format.Merge(f.read(), graph_def)
+ else:
+ from tensorflow.python.platform import gfile
with gfile.FastGFile(get_path(path), "rb") as f:
graph_def.ParseFromString(f.read())
