def testMultipleMetaGraphDef(self):
"""Test saved model with multiple MetaGraphDef."""
saved_model_dir = os.path.join(self.get_temp_dir(), "savedmodel_two_mgd")
builder = saved_model.builder.SavedModelBuilder(saved_model_dir)
with session.Session(graph=ops.Graph()) as sess:
# MetaGraphDef 1
in_tensor = array_ops.placeholder(shape=[1, 28, 28], dtype=dtypes.float32)
out_tensor = in_tensor + in_tensor
sig_input_tensor = saved_model.utils.build_tensor_info(in_tensor)
sig_input_tensor_signature = {"x": sig_input_tensor}
sig_output_tensor = saved_model.utils.build_tensor_info(out_tensor)
sig_output_tensor_signature = {"y": sig_output_tensor}
predict_signature_def = (
saved_model.signature_def_utils.build_signature_def(
sig_input_tensor_signature, sig_output_tensor_signature,
saved_model.signature_constants.PREDICT_METHOD_NAME))
signature_def_map = {
saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
predict_signature_def
}
builder.add_meta_graph_and_variables(
sess,
tags=[saved_model.tag_constants.SERVING, "additional_test_tag"],
signature_def_map=signature_def_map)
# MetaGraphDef 2
builder.add_meta_graph(tags=["tflite"])
builder.save(True)
# Convert to tflite
convert_saved_model.convert(
saved_model_dir=saved_model_dir,
tag_set=set([saved_model.tag_constants.SERVING, "additional_test_tag"]))
def testSimpleSavedModelWithMoreNoneInShape(self):
"""Test a simple savedmodel, fail as more None in input shape."""
saved_model_dir = self._createSimpleSavedModel(
shape=[None, 16, None, 3])
# Convert to tflite: this should raise ValueError, as 3rd dim is None.
with self.assertRaises(ValueError):
convert_saved_model.convert(saved_model_dir=saved_model_dir)
def testSimpleSavedModelWithWrongSignatureKey(self):
"""Test a simple savedmodel, fail as given signature is invalid."""
saved_model_dir = self._createSimpleSavedModel(shape=[1, 16, 16, 3])
# Convert to tflite: this should raise ValueError, as
# signature_key does not exit in the saved_model.
with self.assertRaises(ValueError):
convert_saved_model.convert(
saved_model_dir=saved_model_dir, signature_key="wrong-key")
def testSimpleSavedModelWithWrongSignatureKey(self):
"""Test a simple savedmodel, fail as given signature is invalid."""
saved_model_dir = self._createSimpleSavedModel(shape=[1, 16, 16, 3])
# Convert to tflite: this should raise ValueError, as
# signature_key does not exit in the saved_model.
with self.assertRaises(ValueError):
convert_saved_model.convert(saved_model_dir=saved_model_dir,
signature_key="wrong-key")
def testSimpleSavedModelWithWrongOutputArray(self):
"""Test a simple savedmodel, fail as given output_arrays is invalid."""
# Create a simple savedmodel
saved_model_dir = self._createSimpleSavedModel(shape=[1, 16, 16, 3])
# Convert to tflite: this should raise ValueError, as
# output_arrays is not valid for the saved_model.
with self.assertRaises(ValueError):
convert_saved_model.convert(
saved_model_dir=saved_model_dir, output_arrays="wrong-output")
def testSimpleSavedModelWithWrongOutputArray(self):
"""Test a simple savedmodel, fail as given output_arrays is invalid."""
# Create a simple savedmodel
saved_model_dir = self._createSimpleSavedModel(shape=[1, 16, 16, 3])
# Convert to tflite: this should raise ValueError, as
# output_arrays is not valid for the saved_model.
with self.assertRaises(ValueError):
convert_saved_model.convert(saved_model_dir=saved_model_dir,
output_arrays="wrong-output")
def testSimpleSavedModel(self): """Test a simple savedmodel created on the fly.""" # Create a simple savedmodel saved_model_dir = self._createSimpleSavedModel(shape=[1, 16, 16, 3]) # Convert to tflite result = convert_saved_model.convert(saved_model_dir=saved_model_dir) self.assertTrue(result)
def testTrainedMnistSavedModel(self):
"""Test mnist savedmodel, trained with dummy data and small steps."""
# Build classifier
classifier = estimator.Estimator(
model_fn=model_fn,
params={
"data_format": "channels_last" # tflite format
})
# Train and pred for serving
classifier.train(input_fn=dummy_input_fn, steps=2)
image = array_ops.placeholder(dtypes.float32, [None, 28, 28])
pred_input_fn = estimator.export.build_raw_serving_input_receiver_fn({
"image": image,
})
# Export savedmodel
saved_model_dir = os.path.join(self.get_temp_dir(), "mnist_savedmodel")
classifier.export_savedmodel(saved_model_dir, pred_input_fn)
# Convert to tflite and test output
saved_model_name = os.listdir(saved_model_dir)[0]
saved_model_final_dir = os.path.join(saved_model_dir, saved_model_name)
output_tflite = os.path.join(saved_model_dir,
saved_model_final_dir + ".lite")
# TODO(zhixianyan): no need to limit output_arrays to `Softmax'
# once b/74205001 fixed and argmax implemented in tflite.
result = convert_saved_model.convert(
saved_model_dir=saved_model_final_dir,
output_arrays="Softmax",
output_tflite=output_tflite)
self.assertTrue(result)
def testSimpleSavedModel(self):
"""Test a simple savedmodel created on the fly."""
# Create a simple savedmodel
saved_model_dir = self._createSimpleSavedModel(shape=[1, 16, 16, 3])
# Convert to tflite
result = convert_saved_model.convert(saved_model_dir=saved_model_dir)
self.assertTrue(result)
def testTrainedMnistSavedModel(self):
"""Test mnist savedmodel, trained with dummy data and small steps."""
# Build classifier
classifier = estimator.Estimator(
model_fn=model_fn,
params={
"data_format": "channels_last" # tflite format
})
# Train and pred for serving
classifier.train(input_fn=dummy_input_fn, steps=2)
image = array_ops.placeholder(dtypes.float32, [None, 28, 28])
pred_input_fn = estimator.export.build_raw_serving_input_receiver_fn({
"image":
image,
})
# Export savedmodel
saved_model_dir = os.path.join(self.get_temp_dir(), "mnist_savedmodel")
classifier.export_savedmodel(saved_model_dir, pred_input_fn)
# Convert to tflite and test output
saved_model_name = os.listdir(saved_model_dir)[0]
saved_model_final_dir = os.path.join(saved_model_dir, saved_model_name)
output_tflite = os.path.join(saved_model_dir,
saved_model_final_dir + ".lite")
# TODO(zhixianyan): no need to limit output_arrays to `Softmax'
# once b/74205001 fixed and argmax implemented in tflite.
result = convert_saved_model.convert(
saved_model_dir=saved_model_final_dir,
output_arrays="Softmax",
output_tflite=output_tflite)
self.assertTrue(result)
def testMultipleMetaGraphDef(self):
"""Test saved model with multiple MetaGraphDef."""
saved_model_dir = os.path.join(self.get_temp_dir(),
"savedmodel_two_mgd")
builder = saved_model.builder.SavedModelBuilder(saved_model_dir)
with session.Session(graph=ops.Graph()) as sess:
# MetaGraphDef 1
in_tensor = array_ops.placeholder(shape=[1, 28, 28],
dtype=dtypes.float32)
out_tensor = in_tensor + in_tensor
sig_input_tensor = saved_model.utils.build_tensor_info(in_tensor)
sig_input_tensor_signature = {"x": sig_input_tensor}
sig_output_tensor = saved_model.utils.build_tensor_info(out_tensor)
sig_output_tensor_signature = {"y": sig_output_tensor}
predict_signature_def = (
saved_model.signature_def_utils.build_signature_def(
sig_input_tensor_signature, sig_output_tensor_signature,
saved_model.signature_constants.PREDICT_METHOD_NAME))
signature_def_map = {
saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
predict_signature_def
}
builder.add_meta_graph_and_variables(
sess,
tags=[
saved_model.tag_constants.SERVING, "additional_test_tag"
],
signature_def_map=signature_def_map)
# MetaGraphDef 2
builder.add_meta_graph(tags=["tflite"])
builder.save(True)
# Convert to tflite
convert_saved_model.convert(saved_model_dir=saved_model_dir,
tag_set=set([
saved_model.tag_constants.SERVING,
"additional_test_tag"
]))
def testSimpleSavedModelWithMoreNoneInShape(self):
"""Test a simple savedmodel, fail as more None in input shape."""
saved_model_dir = self._createSimpleSavedModel(shape=[None, 16, None, 3])
# Convert to tflite: this should raise ValueError, as 3rd dim is None.
with self.assertRaises(ValueError):
convert_saved_model.convert(saved_model_dir=saved_model_dir)
def testSimpleSavedModelWithNoneBatchSizeInShape(self): """Test a simple savedmodel, with None in input tensor's shape.""" saved_model_dir = self._createSimpleSavedModel(shape=[None, 16, 16, 3]) result = convert_saved_model.convert(saved_model_dir=saved_model_dir) self.assertTrue(result)
def testSimpleSavedModelWithNoneBatchSizeInShape(self):
"""Test a simple savedmodel, with None in input tensor's shape."""
saved_model_dir = self._createSimpleSavedModel(shape=[None, 16, 16, 3])
result = convert_saved_model.convert(saved_model_dir=saved_model_dir)
self.assertTrue(result)