def testSetTensorShapeEmpty(self):
tensor = array_ops.placeholder(shape=[None, 3, 5],
dtype=dtypes.float32)
self.assertEqual([None, 3, 5], tensor.shape.as_list())
convert_saved_model.set_tensor_shapes([tensor], {})
self.assertEqual([None, 3, 5], tensor.shape.as_list())
def testSetTensorShapeInvalid(self):
tensor = array_ops.placeholder(shape=[None, 3, 5], dtype=dtypes.float32)
self.assertEqual([None, 3, 5], tensor.shape.as_list())
convert_saved_model.set_tensor_shapes([tensor],
{"invalid-input": [5, 3, 5]})
self.assertEqual([None, 3, 5], tensor.shape.as_list())
def testSetTensorShapeNoneValid(self):
tensor = array_ops.placeholder(dtype=dtypes.float32)
self.assertEqual(None, tensor.shape)
convert_saved_model.set_tensor_shapes([tensor],
{"Placeholder": [1, 3, 5]})
self.assertEqual([1, 3, 5], tensor.shape.as_list())
def from_keras_model_file(cls,
model_file,
input_arrays=None,
input_shapes=None,
output_arrays=None):
"""Creates a TocoConverter class from a tf.keras model file.
Args:
model_file: Full filepath of HDF5 file containing the tf.keras model.
input_arrays: List of input tensors to freeze graph with. Uses input
arrays from SignatureDef when none are provided. (default None)
input_shapes: Dict of strings representing input tensor names to list of
integers representing input shapes (e.g., {"foo" : [1, 16, 16, 3]}).
Automatically determined when input shapes is None (e.g., {"foo" :
None}). (default None)
output_arrays: List of output tensors to freeze graph with. Uses output
arrays from SignatureDef when none are provided. (default None)
Returns:
TocoConverter class.
"""
_keras.backend.clear_session()
_keras.backend.set_learning_phase(False)
keras_model = _keras.models.load_model(model_file)
sess = _keras.backend.get_session()
# Get input and output tensors.
if input_arrays:
input_tensors = get_tensors_from_tensor_names(
sess.graph, input_arrays)
else:
input_tensors = keras_model.inputs
if output_arrays:
output_tensors = get_tensors_from_tensor_names(
sess.graph, output_arrays)
else:
output_tensors = keras_model.outputs
set_tensor_shapes(input_tensors, input_shapes)
graph_def = _freeze_graph(sess, output_tensors)
return cls(graph_def, input_tensors, output_tensors)
def from_keras_model_file(cls,
model_file,
input_arrays=None,
input_shapes=None,
output_arrays=None):
"""Creates a TocoConverter class from a tf.keras model file.
Args:
model_file: Full filepath of HDF5 file containing the tf.keras model.
input_arrays: List of input tensors to freeze graph with. Uses input
arrays from SignatureDef when none are provided. (default None)
input_shapes: Dict of strings representing input tensor names to list of
integers representing input shapes (e.g., {"foo" : [1, 16, 16, 3]}).
Automatically determined when input shapes is None (e.g., {"foo" :
None}). (default None)
output_arrays: List of output tensors to freeze graph with. Uses output
arrays from SignatureDef when none are provided. (default None)
Returns:
TocoConverter class.
"""
_keras.backend.clear_session()
_keras.backend.set_learning_phase(False)
keras_model = _keras.models.load_model(model_file)
sess = _keras.backend.get_session()
# Get input and output tensors.
if input_arrays:
input_tensors = get_tensors_from_tensor_names(sess.graph, input_arrays)
else:
input_tensors = keras_model.inputs
if output_arrays:
output_tensors = get_tensors_from_tensor_names(sess.graph, output_arrays)
else:
output_tensors = keras_model.outputs
set_tensor_shapes(input_tensors, input_shapes)
graph_def = _freeze_graph(sess, output_tensors)
return cls(graph_def, input_tensors, output_tensors)
def from_frozen_graph(cls,
graph_def_file,
input_arrays,
output_arrays,
input_shapes=None):
"""Creates a TocoConverter class from a file containing a frozen GraphDef.
Args:
graph_def_file: Full filepath of file containing TensorFlow GraphDef.
input_arrays: List of input tensors to freeze graph with.
output_arrays: List of output tensors to freeze graph with.
input_shapes: Dict of strings representing input tensor names to list of
integers representing input shapes (e.g., {"foo" : [1, 16, 16, 3]}).
Automatically determined when input shapes is None (e.g., {"foo" :
None}). (default None)
Returns:
TocoConverter class.
Raises:
ValueError:
Unable to parse input file.
The graph is not frozen.
input_arrays or output_arrays contains an invalid tensor name.
"""
with _session.Session() as sess:
sess.run(global_variables_initializer())
# Read GraphDef from file.
graph_def = _graph_pb2.GraphDef()
with open(graph_def_file, "rb") as f:
file_content = f.read()
try:
graph_def.ParseFromString(file_content)
except (_text_format.ParseError, DecodeError):
try:
print("Ignore 'tcmalloc: large alloc' warnings.")
if not isinstance(file_content, str):
if PY3:
file_content = file_content.decode('utf-8')
else:
file_content = file_content.encode('utf-8')
_text_format.Merge(file_content, graph_def)
except (_text_format.ParseError, DecodeError):
raise ValueError("Unable to parse input file '{}'.".format(
graph_def_file))
sess.graph.as_default()
import_graph_def(graph_def, name="")
# Get input and output tensors.
input_tensors = get_tensors_from_tensor_names(
sess.graph, input_arrays)
output_tensors = get_tensors_from_tensor_names(
sess.graph, output_arrays)
set_tensor_shapes(input_tensors, input_shapes)
# Check if graph is frozen.
if not _is_frozen_graph(sess):
raise ValueError(
"Please freeze the graph using freeze_graph.py")
# Create TocoConverter class.
return cls(sess.graph_def, input_tensors, output_tensors)
def testSetTensorShapeNoneValid(self):
tensor = array_ops.placeholder(dtype=dtypes.float32)
self.assertEqual(None, tensor.shape)
convert_saved_model.set_tensor_shapes([tensor], {"Placeholder": [1, 3, 5]})
self.assertEqual([1, 3, 5], tensor.shape.as_list())
def from_frozen_graph(cls,
graph_def_file,
input_arrays,
output_arrays,
input_shapes=None):
"""Creates a TocoConverter class from a file containing a frozen GraphDef.
Args:
graph_def_file: Full filepath of file containing TensorFlow GraphDef.
input_arrays: List of input tensors to freeze graph with.
output_arrays: List of output tensors to freeze graph with.
input_shapes: Dict of strings representing input tensor names to list of
integers representing input shapes (e.g., {"foo" : [1, 16, 16, 3]}).
Automatically determined when input shapes is None (e.g., {"foo" :
None}). (default None)
Returns:
TocoConverter class.
Raises:
ValueError:
Unable to parse input file.
The graph is not frozen.
input_arrays or output_arrays contains an invalid tensor name.
"""
with _session.Session() as sess:
sess.run(global_variables_initializer())
# Read GraphDef from file.
graph_def = _graph_pb2.GraphDef()
with open(graph_def_file, "rb") as f:
file_content = f.read()
try:
graph_def.ParseFromString(file_content)
except (_text_format.ParseError, DecodeError):
try:
print("Ignore 'tcmalloc: large alloc' warnings.")
if not isinstance(file_content, str):
if PY3:
file_content = file_content.decode('utf-8')
else:
file_content = file_content.encode('utf-8')
_text_format.Merge(file_content, graph_def)
except (_text_format.ParseError, DecodeError):
raise ValueError(
"Unable to parse input file '{}'.".format(graph_def_file))
sess.graph.as_default()
import_graph_def(graph_def, name="")
# Get input and output tensors.
input_tensors = get_tensors_from_tensor_names(sess.graph, input_arrays)
output_tensors = get_tensors_from_tensor_names(sess.graph, output_arrays)
set_tensor_shapes(input_tensors, input_shapes)
# Check if graph is frozen.
if not _is_frozen_graph(sess):
raise ValueError("Please freeze the graph using freeze_graph.py.")
# Create TocoConverter class.
return cls(sess.graph_def, input_tensors, output_tensors)
