wehub-resource-sync
496 lines
18 KiB
Python
496 lines
18 KiB
Python
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ==============================================================================
|
|
"""TensorFlow Lite Python Interface: Sanity check."""
|
|
from unittest import mock
|
|
|
|
import numpy as np
|
|
|
|
from tensorflow.compiler.mlir.lite import converter_flags_pb2 as _conversion_flags_pb2
|
|
from tensorflow.compiler.mlir.lite.metrics import converter_error_data_pb2
|
|
from tensorflow.lite.python import convert
|
|
from tensorflow.lite.python import op_hint
|
|
from tensorflow.lite.python.interpreter import Interpreter
|
|
from tensorflow.lite.python.metrics.wrapper import metrics_wrapper
|
|
from tensorflow.python.client import session
|
|
from tensorflow.python.framework import dtypes
|
|
from tensorflow.python.framework import ops
|
|
from tensorflow.python.framework import test_util
|
|
from tensorflow.python.framework.graph_util_impl import _bfs_for_reachable_nodes
|
|
from tensorflow.python.framework.graph_util_impl import _extract_graph_summary
|
|
from tensorflow.python.framework.graph_util_impl import _node_name
|
|
from tensorflow.python.ops import array_ops
|
|
from tensorflow.python.ops import array_ops_stack
|
|
from tensorflow.python.ops import math_ops
|
|
from tensorflow.python.platform import test
|
|
|
|
|
|
def _mock_wrapped_convert(
|
|
unused_model_flags_str="",
|
|
conversion_flags_str="",
|
|
unused_input_data_str="",
|
|
unused_debug_info_str="",
|
|
):
|
|
# Simulate the converter throwing and error when
|
|
# `guarantee_all_funcs_one_use` is not set.
|
|
if not _conversion_flags_pb2.ConverterFlags.FromString(
|
|
conversion_flags_str
|
|
).guarantee_all_funcs_one_use:
|
|
raise Exception()
|
|
else:
|
|
return bytes("A model", encoding="utf-8")
|
|
|
|
|
|
def _mock_retrieve_errors():
|
|
err_data = converter_error_data_pb2.ConverterErrorData(
|
|
error_code=converter_error_data_pb2.ConverterErrorData.ERROR_STATEFUL_PARTITIONED_CALL_IN_FINAL_IR
|
|
)
|
|
return [err_data]
|
|
|
|
|
|
class ConvertTest(test_util.TensorFlowTestCase):
|
|
|
|
def testBasic(self):
|
|
with ops.Graph().as_default():
|
|
in_tensor = array_ops.placeholder(
|
|
shape=[1, 16, 16, 3], dtype=dtypes.float32
|
|
)
|
|
out_tensor = in_tensor + in_tensor
|
|
sess = session.Session()
|
|
|
|
# Try running on valid graph
|
|
tflite_model = convert.convert_graphdef(
|
|
sess.graph_def, input_tensors=[in_tensor], output_tensors=[out_tensor]
|
|
)
|
|
self.assertTrue(tflite_model)
|
|
|
|
@mock.patch.object(
|
|
convert.wrap_converter, "wrapped_convert", new=_mock_wrapped_convert
|
|
)
|
|
@mock.patch.object(
|
|
metrics_wrapper, "retrieve_collected_errors", new=_mock_retrieve_errors
|
|
)
|
|
# This test wants to check that in the case of the converter throwing an
|
|
# `ERROR_STATEFUL_PARTITIONED_CALL_IN_FINAL_IR` error, it will
|
|
# retry conversion with the `guarantee_all_funcs_one_use` flag.
|
|
# We can wrap the convert call in order to assert it is called appropriately.
|
|
@mock.patch.object(convert, "convert", wraps=convert.convert)
|
|
def testConversionStatefulPartitionRetry(self, mock_convert):
|
|
with ops.Graph().as_default():
|
|
in_tensor = array_ops.placeholder(
|
|
shape=[1, 16, 16, 3], dtype=dtypes.float32
|
|
)
|
|
out_tensor = in_tensor + in_tensor
|
|
sess = session.Session()
|
|
|
|
model = convert.convert_graphdef(
|
|
sess.graph_def,
|
|
input_tensors=[in_tensor],
|
|
output_tensors=[out_tensor],
|
|
guarantee_all_funcs_one_use=False,
|
|
)
|
|
self.assertTrue(str(model, encoding="utf-8"), "A model")
|
|
self.assertEqual(mock_convert.call_count, 2)
|
|
|
|
def testQuantization(self):
|
|
with ops.Graph().as_default():
|
|
in_tensor = array_ops.placeholder(
|
|
shape=[1, 16, 16, 3], dtype=dtypes.float32
|
|
)
|
|
out_tensor = array_ops.fake_quant_with_min_max_args(
|
|
in_tensor + in_tensor, min=0.0, max=1.0
|
|
)
|
|
sess = session.Session()
|
|
|
|
tflite_model = convert.convert_graphdef(
|
|
sess.graph_def,
|
|
input_tensors=[in_tensor],
|
|
output_tensors=[out_tensor],
|
|
inference_type=dtypes.uint8,
|
|
quantized_input_stats=[(0.0, 1.0)],
|
|
)
|
|
self.assertTrue(tflite_model)
|
|
|
|
def testGraphDefBasic(self):
|
|
with ops.Graph().as_default():
|
|
in_tensor = array_ops.placeholder(
|
|
shape=[1, 16, 16, 3], dtype=dtypes.float32, name="input"
|
|
)
|
|
_ = in_tensor + in_tensor
|
|
sess = session.Session()
|
|
|
|
tflite_model = convert.convert_graphdef_with_arrays(
|
|
sess.graph_def,
|
|
input_arrays_with_shape=[("input", [1, 16, 16, 3])],
|
|
output_arrays=["add"],
|
|
control_output_arrays=None,
|
|
inference_type=dtypes.float32,
|
|
)
|
|
self.assertTrue(tflite_model)
|
|
|
|
# Check values from converted model.
|
|
interpreter = Interpreter(model_content=tflite_model)
|
|
interpreter.allocate_tensors()
|
|
|
|
input_details = interpreter.get_input_details()
|
|
self.assertEqual(1, len(input_details))
|
|
self.assertEqual("input", input_details[0]["name"])
|
|
self.assertEqual(np.float32, input_details[0]["dtype"])
|
|
self.assertTrue(([1, 16, 16, 3] == input_details[0]["shape"]).all()) # type: ignore
|
|
self.assertEqual((0.0, 0.0), input_details[0]["quantization"])
|
|
|
|
output_details = interpreter.get_output_details()
|
|
self.assertEqual(1, len(output_details))
|
|
self.assertEqual("add", output_details[0]["name"])
|
|
self.assertEqual(np.float32, output_details[0]["dtype"])
|
|
self.assertTrue(([1, 16, 16, 3] == output_details[0]["shape"]).all()) # type: ignore
|
|
self.assertEqual((0.0, 0.0), output_details[0]["quantization"])
|
|
|
|
def testGraphDefQuantization(self):
|
|
with ops.Graph().as_default():
|
|
in_tensor_1 = array_ops.placeholder(
|
|
shape=[1, 16, 16, 3], dtype=dtypes.float32, name="inputA"
|
|
)
|
|
in_tensor_2 = array_ops.placeholder(
|
|
shape=[1, 16, 16, 3], dtype=dtypes.float32, name="inputB"
|
|
)
|
|
_ = array_ops.fake_quant_with_min_max_args(
|
|
in_tensor_1 + in_tensor_2, min=0.0, max=1.0, name="output"
|
|
)
|
|
sess = session.Session()
|
|
|
|
tflite_model = convert.convert_graphdef_with_arrays(
|
|
sess.graph_def,
|
|
input_arrays_with_shape=[
|
|
("inputA", [1, 16, 16, 3]),
|
|
("inputB", [1, 16, 16, 3]),
|
|
],
|
|
output_arrays=["output"],
|
|
control_output_arrays=None,
|
|
inference_type=dtypes.uint8,
|
|
quantized_input_stats=[(0.0, 1.0), (0.0, 1.0)],
|
|
)
|
|
self.assertTrue(tflite_model)
|
|
|
|
# Check values from converted model.
|
|
interpreter = Interpreter(model_content=tflite_model)
|
|
interpreter.allocate_tensors()
|
|
|
|
input_details = interpreter.get_input_details()
|
|
self.assertEqual(2, len(input_details))
|
|
self.assertEqual("inputA", input_details[0]["name"])
|
|
self.assertEqual(np.uint8, input_details[0]["dtype"])
|
|
self.assertTrue(([1, 16, 16, 3] == input_details[0]["shape"]).all()) # type: ignore
|
|
self.assertEqual(
|
|
(1.0, 0.0), input_details[0]["quantization"]
|
|
) # scale, zero_point
|
|
|
|
self.assertEqual("inputB", input_details[1]["name"])
|
|
self.assertEqual(np.uint8, input_details[1]["dtype"])
|
|
self.assertTrue(([1, 16, 16, 3] == input_details[1]["shape"]).all()) # type: ignore
|
|
self.assertEqual(
|
|
(1.0, 0.0), input_details[1]["quantization"]
|
|
) # scale, zero_point
|
|
|
|
output_details = interpreter.get_output_details()
|
|
self.assertEqual(1, len(output_details))
|
|
self.assertEqual("output", output_details[0]["name"])
|
|
self.assertEqual(np.uint8, output_details[0]["dtype"])
|
|
self.assertTrue(([1, 16, 16, 3] == output_details[0]["shape"]).all()) # type: ignore
|
|
self.assertGreater(output_details[0]["quantization"][0], 0) # scale
|
|
|
|
def testGraphDefQuantizationInvalid(self):
|
|
with ops.Graph().as_default():
|
|
in_tensor_1 = array_ops.placeholder(
|
|
shape=[1, 16, 16, 3], dtype=dtypes.float32, name="inputA"
|
|
)
|
|
in_tensor_2 = array_ops.placeholder(
|
|
shape=[1, 16, 16, 3], dtype=dtypes.float32, name="inputB"
|
|
)
|
|
_ = array_ops.fake_quant_with_min_max_args(
|
|
in_tensor_1 + in_tensor_2, min=0.0, max=1.0, name="output"
|
|
)
|
|
sess = session.Session()
|
|
|
|
with self.assertRaises(ValueError) as error:
|
|
convert.convert_graphdef_with_arrays(
|
|
sess.graph_def,
|
|
input_arrays_with_shape=[
|
|
("inputA", [1, 16, 16, 3]),
|
|
("inputB", [1, 16, 16, 3]),
|
|
],
|
|
output_arrays=["output"],
|
|
control_output_arrays=None,
|
|
inference_type=dtypes.uint8,
|
|
)
|
|
self.assertEqual(
|
|
"The `quantized_input_stats` flag must be defined when either "
|
|
"`inference_type` flag or `inference_input_type` flag is set to "
|
|
"tf.int8 or tf.uint8.",
|
|
str(error.exception),
|
|
)
|
|
|
|
|
|
class ConvertTestOpHint(test_util.TensorFlowTestCase):
|
|
"""Test the hint to stub functionality."""
|
|
|
|
def _getGraphOpTypes(self, graphdef, output_nodes):
|
|
"""Returns used op types in `graphdef` reachable from `output_nodes`.
|
|
|
|
This is used to check that after the stub transformation the expected
|
|
nodes are there.
|
|
|
|
NOTE: this is not a exact test that the graph is the correct output, but
|
|
it balances compact expressibility of test with sanity checking.
|
|
|
|
Args:
|
|
graphdef: TensorFlow proto graphdef.
|
|
output_nodes: A list of output node names that we need to reach.
|
|
|
|
Returns:
|
|
A set of node types reachable from `output_nodes`.
|
|
"""
|
|
name_to_input_name, name_to_node, _ = _extract_graph_summary(graphdef)
|
|
# Find all nodes that are needed by the outputs
|
|
used_node_names = _bfs_for_reachable_nodes(output_nodes, name_to_input_name)
|
|
return set([name_to_node[node_name].op for node_name in used_node_names])
|
|
|
|
def _countIdentities(self, nodes):
|
|
"""Count the number of "Identity" op types in the list of proto nodes.
|
|
|
|
Args:
|
|
nodes: NodeDefs of the graph.
|
|
|
|
Returns:
|
|
The number of nodes with op type "Identity" found.
|
|
"""
|
|
return len([x for x in nodes if x.op == "Identity"])
|
|
|
|
def testSwishLiteHint(self):
|
|
"""Makes a custom op swish and makes sure it gets converted as a unit."""
|
|
with ops.Graph().as_default():
|
|
image = array_ops.constant([1.0, 2.0, 3.0, 4.0])
|
|
swish_scale = array_ops.constant(1.0)
|
|
|
|
def _swish(input_tensor, scale):
|
|
custom = op_hint.OpHint("cool_activation")
|
|
input_tensor, scale = custom.add_inputs(input_tensor, scale)
|
|
output = math_ops.sigmoid(input_tensor) * input_tensor * scale
|
|
(output,) = custom.add_outputs(output)
|
|
return output
|
|
|
|
output = array_ops.identity(
|
|
_swish(image, swish_scale), name="ModelOutput"
|
|
)
|
|
|
|
with self.cached_session() as sess:
|
|
# check if identities have been put into the graph (2 input, 1 output,
|
|
# and 1 final output).
|
|
self.assertEqual(self._countIdentities(sess.graph_def.node), 4)
|
|
|
|
stubbed_graphdef = op_hint.convert_op_hints_to_stubs(
|
|
graph_def=sess.graph_def
|
|
)
|
|
|
|
self.assertEqual(
|
|
self._getGraphOpTypes(
|
|
stubbed_graphdef,
|
|
output_nodes=[op_hint._tensor_name_base(output.name)],
|
|
),
|
|
set(["cool_activation", "Const", "Identity"]),
|
|
)
|
|
|
|
def testScaleAndBiasAndIdentity(self):
|
|
"""This tests a scaled add which has 3 inputs and 2 outputs."""
|
|
with ops.Graph().as_default():
|
|
a = array_ops.constant(1.0)
|
|
x = array_ops.constant([2.0, 3.0])
|
|
b = array_ops.constant([4.0, 5.0])
|
|
|
|
def _scaled_and_bias_and_identity(a, x, b):
|
|
custom = op_hint.OpHint("scale_and_bias_and_identity")
|
|
a, x, b = custom.add_inputs(a, x, b)
|
|
return custom.add_outputs(a * x + b, x)
|
|
|
|
output = array_ops.identity(
|
|
_scaled_and_bias_and_identity(a, x, b), name="ModelOutput"
|
|
)
|
|
|
|
with self.cached_session() as sess:
|
|
# make sure one identity for each input (3) and output (2) => 3 + 2 = 5
|
|
# +1 for the final output
|
|
self.assertEqual(self._countIdentities(sess.graph_def.node), 6)
|
|
|
|
stubbed_graphdef = op_hint.convert_op_hints_to_stubs(
|
|
graph_def=sess.graph_def
|
|
)
|
|
|
|
self.assertEqual(
|
|
self._getGraphOpTypes(
|
|
stubbed_graphdef,
|
|
output_nodes=[op_hint._tensor_name_base(output.name)],
|
|
),
|
|
set(["scale_and_bias_and_identity", "Const", "Identity", "Pack"]),
|
|
)
|
|
|
|
def testTwoFunctions(self):
|
|
"""Tests if two functions are converted correctly."""
|
|
with ops.Graph().as_default():
|
|
a = array_ops.constant([1.0])
|
|
b = array_ops.constant([1.0])
|
|
|
|
def _double_values(x):
|
|
custom = op_hint.OpHint("add_test")
|
|
(x,) = custom.add_inputs(x)
|
|
output = math_ops.multiply(x, x)
|
|
(output,) = custom.add_outputs(output)
|
|
return output
|
|
|
|
output = array_ops.identity(
|
|
math_ops.add(_double_values(a), _double_values(b)), name="ModelOutput"
|
|
)
|
|
|
|
with self.cached_session() as sess:
|
|
# make sure one identity for each input (2) and output (2) => 2 + 2
|
|
# +1 for the final output
|
|
self.assertEqual(self._countIdentities(sess.graph_def.node), 5)
|
|
stubbed_graphdef = op_hint.convert_op_hints_to_stubs(
|
|
graph_def=sess.graph_def
|
|
)
|
|
self.assertEqual(
|
|
self._getGraphOpTypes(
|
|
stubbed_graphdef,
|
|
output_nodes=[op_hint._tensor_name_base(output.name)],
|
|
),
|
|
set(["add_test", "Const", "Identity", "AddV2"]),
|
|
)
|
|
|
|
def _get_input_index(self, x):
|
|
return x.op.node_def.attr[op_hint.OpHint.FUNCTION_INPUT_INDEX_ATTR].i
|
|
|
|
def _get_output_index(self, x):
|
|
return x.op.node_def.attr[op_hint.OpHint.FUNCTION_OUTPUT_INDEX_ATTR].i
|
|
|
|
def _get_sort_index(self, x):
|
|
return x.op.node_def.attr[op_hint.OpHint.FUNCTION_SORT_INDEX_ATTR].i
|
|
|
|
def testTags(self):
|
|
"""Test if multiple args with the same tag are grouped."""
|
|
with ops.Graph().as_default():
|
|
a = array_ops.constant([1.0])
|
|
b = array_ops.constant([2.0])
|
|
c = array_ops.constant([3.0])
|
|
d = array_ops.constant([4.0])
|
|
custom = op_hint.OpHint("test_tag")
|
|
a = custom.add_input(
|
|
a, tag="mytag", aggregate=op_hint.OpHint.AGGREGATE_STACK
|
|
)
|
|
(b,) = custom.add_inputs(b)
|
|
c = custom.add_input(
|
|
c, tag="mytag", aggregate=op_hint.OpHint.AGGREGATE_STACK
|
|
)
|
|
d = custom.add_input(
|
|
d, tag="mytag2", aggregate=op_hint.OpHint.AGGREGATE_STACK
|
|
)
|
|
res = math_ops.add(math_ops.mul(a, b), math_ops.mul(c, b))
|
|
custom.add_outputs([res])
|
|
with self.cached_session():
|
|
self.assertEqual(self._get_input_index(a), 0)
|
|
self.assertEqual(self._get_sort_index(a), 0)
|
|
self.assertEqual(self._get_input_index(b), 1)
|
|
self.assertEqual(self._get_sort_index(b), 0)
|
|
self.assertEqual(self._get_input_index(c), 0)
|
|
self.assertEqual(self._get_sort_index(c), 1)
|
|
|
|
def testOverrideIndex(self):
|
|
with ops.Graph().as_default():
|
|
a = array_ops.constant([1.0])
|
|
b = array_ops.constant([2.0])
|
|
c = array_ops.constant([3.0])
|
|
custom = op_hint.OpHint("test_override")
|
|
b = custom.add_input(b) # should auto assign 0
|
|
a = custom.add_input(a, index_override=1)
|
|
c = custom.add_input(c) # should auto assign 2
|
|
with self.cached_session():
|
|
self.assertEqual(self._get_input_index(a), 1)
|
|
self.assertEqual(self._get_input_index(b), 0)
|
|
self.assertEqual(self._get_input_index(c), 2)
|
|
|
|
def testAggregate(self):
|
|
with ops.Graph().as_default():
|
|
a = array_ops.constant([3.0, 4.0])
|
|
b = array_ops.constant([5.0, 6.0])
|
|
hint = op_hint.OpHint("agg")
|
|
a0, a1 = array_ops_stack.unstack(a)
|
|
b0, b1 = array_ops_stack.unstack(b)
|
|
|
|
a0 = hint.add_input(a0, tag="c", aggregate=op_hint.OpHint.AGGREGATE_STACK)
|
|
b0 = hint.add_input(b0, tag="n", aggregate=op_hint.OpHint.AGGREGATE_STACK)
|
|
a1 = hint.add_input(a1, tag="c", aggregate=op_hint.OpHint.AGGREGATE_STACK)
|
|
b1 = hint.add_input(b1, tag="n", aggregate=op_hint.OpHint.AGGREGATE_STACK)
|
|
|
|
c0 = math_ops.add(a0, b0, name="addleft")
|
|
c1 = math_ops.add(a1, b1, name="addright")
|
|
c0 = hint.add_output(
|
|
c0, tag="out", aggregate=op_hint.OpHint.AGGREGATE_STACK
|
|
)
|
|
c1 = hint.add_output(
|
|
c1, tag="out", aggregate=op_hint.OpHint.AGGREGATE_STACK
|
|
)
|
|
|
|
curr = array_ops_stack.stack([c0, c1])
|
|
output = array_ops.identity(curr, name="FINAL_OUTPUT")
|
|
with self.cached_session() as sess:
|
|
stubbed_graphdef = op_hint.convert_op_hints_to_stubs(
|
|
graph_def=sess.graph_def
|
|
)
|
|
self.assertEqual(
|
|
self._getGraphOpTypes(
|
|
stubbed_graphdef,
|
|
output_nodes=[op_hint._tensor_name_base(output.name)],
|
|
),
|
|
set(["agg", "Const", "Identity"]),
|
|
)
|
|
|
|
def testFindHintedOutputNodes(self):
|
|
"""Test if all hinted output nodes are correctly found."""
|
|
with ops.Graph().as_default():
|
|
|
|
def _build_ophinted_op(name, input1, input2):
|
|
custom_op = op_hint.OpHint(name)
|
|
input1 = custom_op.add_input(input1)
|
|
input2 = custom_op.add_input(input2)
|
|
output = math_ops.mul(input1, input2)
|
|
return custom_op.add_output(output)
|
|
|
|
output_1 = _build_ophinted_op(
|
|
"custom_op_1", array_ops.constant([1.0]), array_ops.constant([2.0])
|
|
)
|
|
output_2 = _build_ophinted_op(
|
|
"custom_op_2", array_ops.constant([3.0]), array_ops.constant([4.0])
|
|
)
|
|
with self.cached_session() as sess:
|
|
hinted_outputs_nodes = op_hint.find_all_hinted_output_nodes(sess)
|
|
expected_hinted_output_nodes = [
|
|
_node_name(output_1.name),
|
|
_node_name(output_2.name),
|
|
]
|
|
self.assertEqual(
|
|
len(hinted_outputs_nodes), len(expected_hinted_output_nodes)
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test.main()
|