Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Reading Tensorflow SavedModels #12

Merged
merged 3 commits into from
Apr 9, 2018
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 9 additions & 5 deletions maraboupy/Marabou.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,15 +14,19 @@ def read_nnet(filename):
return MarabouNetworkNNet(filename)


def read_tf(filename, inputName=None, outputName=None):
def read_tf(filename, inputName=None, outputName=None, savedModel=False, savedModelTags=[]):
"""
Constructs a MarabouNetworkTF object from a frozen Tensorflow protobuf

Args:
filename: (string) path to the .nnet file.
inputName: (string) optional, name of operation corresponding to input
outputName: (string) optional, name of operation corresponding to output
filename: (string) If savedModel is false, path to the frozen graph .pb file.
If savedModel is true, path to SavedModel folder, which
contains .pb file and variables subdirectory.
inputName: (string) optional, name of operation corresponding to input.
outputName: (string) optional, name of operation corresponding to output.
savedModel: (bool) If false, load frozen graph. If true, load SavedModel object.
savedModelTags: (list of strings) If loading a SavedModel, the user must specify tags used.
Returns:
marabouNetworkTF: (MarabouNetworkTF) representing network
"""
return MarabouNetworkTF(filename, inputName, outputName)
return MarabouNetworkTF(filename, inputName, outputName, savedModel, savedModelTags)
72 changes: 50 additions & 22 deletions maraboupy/MarabouNetworkTF.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import numpy as np
from tensorflow.python.framework import tensor_util
from tensorflow.python.framework import graph_util
import os
os.environ['TF_CPP_MIN_LOG_LEVEL']='2'
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
Expand All @@ -9,17 +10,21 @@
from . import MarabouNetwork

class MarabouNetworkTF(MarabouNetwork.MarabouNetwork):
def __init__(self, filename, inputName=None, outputName=None):
def __init__(self, filename, inputName=None, outputName=None, savedModel=False, savedModelTags=[]):
"""
Constructs a MarabouNetworkTF object from a frozen Tensorflow protobuf
Constructs a MarabouNetworkTF object from a frozen Tensorflow protobuf or SavedModel

Args:
filename: (string) path to the .nnet file.
filename: (string) If savedModel is false, path to the frozen graph .pb file.
If savedModel is true, path to SavedModel folder, which
contains .pb file and variables subdirectory.
inputName: (string) optional, name of operation corresponding to input.
outputName: (string) optional, name of operation corresponding to output.
savedModel: (bool) If false, load frozen graph. If true, load SavedModel object.
savedModelTags: (list of strings) If loading a SavedModel, the user must specify tags used.
"""
super().__init__()
self.readFromPb(filename, inputName, outputName)
self.readFromPb(filename, inputName, outputName, savedModel, savedModelTags)

def clear(self):
"""
Expand All @@ -32,23 +37,41 @@ def clear(self):
self.outputOp = None
self.sess = None

def readFromPb(self, filename, inputName, outputName):
def readFromPb(self, filename, inputName, outputName, savedModel, savedModelTags):
"""
Constructs a MarabouNetworkTF object from a frozen Tensorflow protobuf
Constructs a MarabouNetworkTF object from a frozen Tensorflow protobuf or SavedModel

Args:
filename: (string) path to the .pb file.
inputName: (string) optional, name of operation corresponding to input
outputName: (string) optional, name of operation corresponding to output
filename: (string) If savedModel is false, path to the frozen graph .pb file.
If savedModel is true, path to SavedModel folder, which
contains .pb file and variables subdirectory.
inputName: (string) optional, name of operation corresponding to input.
outputName: (string) optional, name of operation corresponding to output.
savedModel: (bool) If false, load frozen graph. If true, load SavedModel object.
savedModelTags: (list of strings) If loading a SavedModel, the user must specify tags used.
"""
### Read protobuf file and begin session ###
with tf.gfile.GFile(filename, "rb") as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
with tf.Graph().as_default() as graph:
tf.import_graph_def(graph_def, name="")
self.sess = tf.Session(graph=graph)
### END reading protobuf ###

if savedModel:
### Read SavedModel ###
sess = tf.Session()
tf.saved_model.loader.load(sess, savedModelTags, filename)

### Simplify graph using outputName, which must be specified for SavedModel ###
simp_graph_def = graph_util.convert_variables_to_constants(sess,sess.graph.as_graph_def(),[outputName])
with tf.Graph().as_default() as graph:
tf.import_graph_def(simp_graph_def, name="")
self.sess = tf.Session(graph=graph)
### End reading SavedModel

else:
### Read protobuf file and begin session ###
with tf.gfile.GFile(filename, "rb") as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
with tf.Graph().as_default() as graph:
tf.import_graph_def(graph_def, name="")
self.sess = tf.Session(graph=graph)
### END reading protobuf ###

### Find operations corresponding to input and output ###
if inputName:
Expand Down Expand Up @@ -78,8 +101,11 @@ def setInputOp(self, op):
Arguments:
op: (tf.op) Representing input
"""
shape = tuple(op.outputs[0].shape.as_list())
self.shapeMap[op] = shape
try:
shape = tuple(op.outputs[0].shape.as_list())
self.shapeMap[op] = shape
except:
self.shapeMap[op] = [None]
self.inputOp = op
self.inputVars = self.opToVarArray(self.inputOp)

Expand All @@ -89,8 +115,11 @@ def setOutputOp(self, op):
Arguments:
op: (tf.op) Representing output
"""
shape = tuple(op.outputs[0].shape.as_list())
self.shapeMap[op] = shape
try:
shape = tuple(op.outputs[0].shape.as_list())
self.shapeMap[op] = shape
except:
self.shapeMap[op] = [None]
self.outputOp = op
self.outputVars = self.opToVarArray(self.outputOp)

Expand All @@ -117,7 +146,6 @@ def opToVarArray(self, x):
### END finding number of new variables ###

v = np.array([self.getNewVariable() for _ in range(size)]).reshape(shape)
self.numVars += size
self.varMap[x] = v
return v

Expand Down