def testAddOp(self, tf_quantization_mode):
root = tracking.AutoTrackable()
root.add_func = def_function.function(lambda x: x + x)
input_data = tf.reshape(tf.range(4, dtype=tf.float32), [1, 4])
concrete_func = root.add_func.get_concrete_function(input_data)
# Convert model and check if the op is not flex.
converter = lite.TFLiteConverterV2.from_concrete_functions([concrete_func],
root)
converter._experimental_tf_quantization_mode = tf_quantization_mode
tflite_model = converter.convert()
self.assertTrue(tflite_model)
if tf_quantization_mode == 'LEGACY_INTEGER':
self.assertIn('ADD', tflite_test_util.get_ops_list(tflite_model))
else:
self.assertIn('FlexAddV2', tflite_test_util.get_ops_list(tflite_model))
# Check the model works.
interpreter = Interpreter(model_content=tflite_model)
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
test_input = np.array([[1.0, 2.0, 3.0, 4.0]], dtype=np.float32)
interpreter.set_tensor(input_details[0]['index'], test_input)
interpreter.invoke()
output_details = interpreter.get_output_details()
expected_output = np.array([[2.0, 4.0, 6.0, 8.0]], dtype=np.float32)
output_data = interpreter.get_tensor(output_details[0]['index'])
self.assertTrue((expected_output == output_data).all())
def testFlexWithDoubleOp(self):
# Create a graph that has one double op.
saved_model_dir = os.path.join(self.get_temp_dir(), 'model2')
with ops.Graph().as_default():
with session.Session() as sess:
in_tensor = array_ops.placeholder(shape=[1, 4],
dtype=dtypes.int32,
name='input')
out_tensor = double_op.double(in_tensor)
inputs = {'x': in_tensor}
outputs = {'z': out_tensor}
saved_model.simple_save(sess, saved_model_dir, inputs, outputs)
converter = lite.TFLiteConverterV2.from_saved_model(saved_model_dir)
converter.target_spec.supported_ops = set([lite.OpsSet.SELECT_TF_OPS])
converter.target_spec.experimental_select_user_tf_ops = ['Double']
tflite_model = converter.convert()
self.assertTrue(tflite_model)
self.assertIn('FlexDouble',
tflite_test_util.get_ops_list(tflite_model))
# Check the model works with TensorFlow ops.
interpreter = Interpreter(model_content=tflite_model)
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
test_input = np.array([[1.0, 2.0, 3.0, 4.0]], dtype=np.int32)
interpreter.set_tensor(input_details[0]['index'], test_input)
interpreter.invoke()
output_details = interpreter.get_output_details()
expected_output = np.array([[2.0, 4.0, 6.0, 8.0]], dtype=np.int32)
output_data = interpreter.get_tensor(output_details[0]['index'])
self.assertTrue((expected_output == output_data).all())
def testL2LossOp(self, tf_quantization_mode):
root = tracking.AutoTrackable()
root.l2_loss_func = def_function.function(lambda x: nn_ops.l2_loss(x)) # pylint: disable=unnecessary-lambda
input_data = tf.range(4, dtype=tf.float32)
concrete_func = root.l2_loss_func.get_concrete_function(input_data)
converter = lite.TFLiteConverterV2.from_concrete_functions([concrete_func],
root)
converter._experimental_tf_quantization_mode = tf_quantization_mode
tflite_model = converter.convert()
self.assertTrue(tflite_model)
self.assertIn('FlexL2Loss', tflite_test_util.get_ops_list(tflite_model))
# Check the model works.
interpreter = Interpreter(model_content=tflite_model)
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
test_input = np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32)
interpreter.set_tensor(input_details[0]['index'], test_input)
interpreter.invoke()
output_details = interpreter.get_output_details()
expected_output = np.array([15.0], dtype=np.float32)
output_data = interpreter.get_tensor(output_details[0]['index'])
self.assertTrue((expected_output == output_data).all())
def _run(self,
flags_str,
should_succeed,
expected_ops_in_converted_model=None,
expected_output_shapes=None):
output_file = os.path.join(self.get_temp_dir(), 'model.tflite')
tflite_bin = resource_loader.get_path_to_datafile('tflite_convert')
cmdline = '{0} --output_file={1} {2}'.format(tflite_bin, output_file,
flags_str)
exitcode = os.system(cmdline)
if exitcode == 0:
with gfile.Open(output_file, 'rb') as model_file:
content = model_file.read()
self.assertEqual(content is not None, should_succeed)
if expected_ops_in_converted_model:
op_set = tflite_test_util.get_ops_list(content)
for opname in expected_ops_in_converted_model:
self.assertIn(opname, op_set)
if expected_output_shapes:
output_shapes = tflite_test_util.get_output_shapes(content)
self.assertEqual(output_shapes, expected_output_shapes)
os.remove(output_file)
else:
self.assertFalse(should_succeed)
def testFlexWithAutomaticPassThrough(self):
# Create a graph that has one L2Loss op.
with ops.Graph().as_default():
with session.Session() as sess:
in_tensor = array_ops.placeholder(
shape=[4], dtype=dtypes.float32, name='input')
out_tensor = nn_ops.l2_loss(in_tensor)
converter = lite.TFLiteConverter.from_session(sess, [in_tensor],
[out_tensor])
converter.target_spec.supported_ops = set([lite.OpsSet.SELECT_TF_OPS])
converter._experimental_allow_all_select_tf_ops = True
tflite_model = converter.convert()
self.assertTrue(tflite_model)
self.assertIn('FlexL2Loss', tflite_test_util.get_ops_list(tflite_model))
def testConvOpWithBias(self, tf_quantization_mode):
class ConvModel(tracking.AutoTrackable):
@def_function.function
def conv_func(self, in_tensor, filter_tensor):
bias = constant_op.constant(3., shape=[1])
conv_tensor = tf.nn.conv2d(
in_tensor,
filter_tensor,
strides=[1, 1, 1, 1],
dilations=[1, 1, 1, 1],
padding='VALID',
data_format='NHWC')
conv_tensor = conv_tensor + bias
return tf.nn.relu(conv_tensor)
root = ConvModel()
input_data = tf.reshape(tf.range(4, dtype=tf.float32), [1, 2, 2, 1])
filter_data = tf.reshape(tf.range(2, dtype=tf.float32), [1, 2, 1, 1])
concrete_func = root.conv_func.get_concrete_function(
input_data, filter_data)
converter = lite.TFLiteConverterV2.from_concrete_functions([concrete_func],
root)
converter._experimental_tf_quantization_mode = tf_quantization_mode
tflite_model = converter.convert()
self.assertTrue(tflite_model)
self.assertCountEqual(['CONV_2D', 'RESHAPE'],
tflite_test_util.get_ops_list(tflite_model))
# Check the model works.
interpreter = Interpreter(model_content=tflite_model)
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
test_input = np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32).reshape(
(1, 2, 2, 1))
interpreter.set_tensor(input_details[0]['index'], test_input)
test_filter = np.array([1.0, 0.0], dtype=np.float32).reshape((1, 2, 1, 1))
interpreter.set_tensor(input_details[1]['index'], test_filter)
interpreter.invoke()
output_details = interpreter.get_output_details()
expected_output = np.array([[[[4.]], [[6.]]]], dtype=np.float32)
output_data = interpreter.get_tensor(output_details[0]['index'])
self.assertTrue((expected_output == output_data).all())
def testFlexWithCustomOp(self):
new_graph, inputs, outputs = self._createGraphWithCustomOp(
opname='CustomAdd4')
# Import to load the custom opdef.
saved_model_dir = os.path.join(self.get_temp_dir(), 'model')
with ops.Graph().as_default():
with session.Session() as sess:
import_graph_def(new_graph, name='')
saved_model.simple_save(sess, saved_model_dir, inputs, outputs)
converter = lite.TFLiteConverterV2.from_saved_model(saved_model_dir)
converter.target_spec.supported_ops = set([lite.OpsSet.SELECT_TF_OPS])
converter.target_spec.experimental_select_user_tf_ops = ['CustomAdd4']
tflite_model = converter.convert()
self.assertIn('FlexCustomAdd4', tflite_test_util.get_ops_list(tflite_model))
def mlir_convert(options, graph_def, input_tensors, output_tensors, **kwargs):
"""Convert a model's graph def into a tflite model with MLIR-based conversion.
Args:
options: A lite.testing.generate_examples_lib.Options instance.
graph_def: A GraphDef object.
input_tensors: List of input tensor tuples `(name, shape, type)`.
output_tensors: List of output tensors (names).
**kwargs: Extra parameters.
Returns:
output tflite model, log_txt from conversion
or None, log_txt if it did not convert properly.
"""
test_params = kwargs.get("test_params", {})
# TODO(b/146025965): Rename ExtraTocoOptions to ExtraConvertOptions or
# something else.
extra_toco_options = kwargs.get("extra_toco_options",
zip_test_utils.ExtraTocoOptions())
input_arrays = [x[0] for x in input_tensors]
input_shapes = zip_test_utils.get_input_shapes_map(input_tensors)
tflite_model = None
log = ""
with tempfile.NamedTemporaryFile() as graphdef_file:
graphdef_file.write(graph_def.SerializeToString())
graphdef_file.flush()
converter = tf.lite.TFLiteConverter.from_frozen_graph(
graphdef_file.name, input_arrays, output_tensors, input_shapes)
converter.allow_custom_ops = extra_toco_options.allow_custom_ops
converter.experimental_new_quantizer = options.mlir_quantizer
if options.run_with_flex:
converter.supported_ops = set(
[tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS])
if test_params.get("dynamic_range_quantize", False):
converter.optimizations = [tf.lite.Optimize.DEFAULT]
if test_params.get("fully_quantize", False):
converter.optimizations = [tf.lite.Optimize.DEFAULT]
# Read the input range for the representative dataset from parameters.
min_value, max_value = test_params.get("input_range", (-1, 1))
def representative_dataset(input_tensors):
calibration_inputs = []
for _, shape, _ in input_tensors:
if shape:
dims = [
1 if dim.value is None else dim.value
for dim in shape.dims
]
calibration_inputs.append(
np.random.uniform(min_value, max_value,
tuple(dims)).astype(np.float32))
return calibration_inputs
def representative_dataset_gen():
for _ in range(100):
yield representative_dataset(input_tensors)
if test_params.get("quant_16x8", False):
converter.target_spec.supported_ops = [
tf.lite.OpsSet.\
EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8
]
else:
converter.target_spec.supported_ops = [
tf.lite.OpsSet.TFLITE_BUILTINS_INT8
]
converter.representative_dataset = representative_dataset_gen
if extra_toco_options.inference_input_type:
converter.inference_input_type = (
extra_toco_options.inference_input_type)
if extra_toco_options.inference_output_type:
converter.inference_output_type = (
extra_toco_options.inference_output_type)
try:
tflite_model = converter.convert()
if options.expected_ops_in_converted_model:
ops_list = tflite_test_util.get_ops_list(tflite_model)
for expected_op in options.expected_ops_in_converted_model:
if expected_op not in ops_list:
# Force the test to fail.
tflite_model = None
raise ValueError(
"{} op not found in the converted model".format(
expected_op))
except Exception as e: # pylint: disable=broad-except
log = str(e)
return tflite_model, log
def mlir_convert(
options,
saved_model_dir,
input_tensors,
output_tensors, # pylint: disable=unused-argument
**kwargs):
"""Convert a saved model into a tflite model with MLIR-based conversion.
Args:
options: A lite.testing.generate_examples_lib.Options instance.
saved_model_dir: Path to the saved model.
input_tensors: List of input tensor tuples `(name, shape, type)`.
output_tensors: List of output tensors (names).
**kwargs: Extra parameters.
Returns:
output tflite model, log_txt from conversion
or None, log_txt if it did not convert properly.
"""
test_params = kwargs.get("test_params", {})
# TODO(b/146025965): Rename ExtraTocoOptions to ExtraConvertOptions or
# something else.
extra_toco_options = kwargs.get("extra_toco_options",
zip_test_utils.ExtraTocoOptions())
tflite_model = None
log = ""
signature_key = signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY
converter = tf.lite.TFLiteConverter.from_saved_model(
saved_model_dir, [signature_key])
converter.allow_custom_ops = extra_toco_options.allow_custom_ops
converter.experimental_new_quantizer = options.mlir_quantizer
if options.make_tf_ptq_tests:
if options.hlo_aware_conversion:
tf_quantization_mode = "DEFAULT"
else:
tf_quantization_mode = "LEGACY_INTEGER"
converter._experimental_tf_quantization_mode = tf_quantization_mode # pylint: disable=protected-access
if options.run_with_flex:
converter.target_spec.supported_ops = set(
[tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS])
if test_params.get("dynamic_range_quantize", False):
converter.optimizations = [tf.lite.Optimize.DEFAULT]
if test_params.get("fully_quantize", False):
converter.optimizations = [tf.lite.Optimize.DEFAULT]
# Read the input range for the representative dataset from parameters.
min_value, max_value = test_params.get("input_range", (-1, 1))
def representative_dataset(input_tensors):
calibration_inputs = {}
for name, shape, dtype in input_tensors:
if shape:
dims = [
1 if dim.value is None else dim.value
for dim in shape.dims
]
calibration_inputs[name] = np.random.uniform(
min_value, max_value,
tuple(dims)).astype(dtype.as_numpy_dtype)
return calibration_inputs
def representative_dataset_gen():
for _ in range(100):
yield representative_dataset(input_tensors)
if test_params.get("quant_16x8", False):
converter.target_spec.supported_ops = [
tf.lite.OpsSet.
EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8
]
else:
converter.target_spec.supported_ops = [
tf.lite.OpsSet.TFLITE_BUILTINS_INT8
]
converter.representative_dataset = representative_dataset_gen
if extra_toco_options.inference_input_type:
converter.inference_input_type = (
extra_toco_options.inference_input_type)
if extra_toco_options.inference_output_type:
converter.inference_output_type = (
extra_toco_options.inference_output_type)
try:
tflite_model = converter.convert()
if options.expected_ops_in_converted_model:
ops_list = tflite_test_util.get_ops_list(tflite_model)
for expected_op in options.expected_ops_in_converted_model:
if expected_op not in ops_list:
# Force the test to fail.
tflite_model = None
raise ValueError(
"{} op not found in the converted model".format(
expected_op))
except Exception as e: # pylint: disable=broad-except
log = str(e)
return tflite_model, log
def testFlexOp(self):
model_path = resource_loader.get_path_to_datafile(
'../testdata/softplus_flex.bin')
op_set = tflite_test_util.get_ops_list(
gfile.GFile(model_path, 'rb').read())
self.assertCountEqual(op_set, ['FlexSoftplus'])
def testBuiltinOp(self):
model_path = resource_loader.get_path_to_datafile(
'../testdata/add.bin')
op_set = tflite_test_util.get_ops_list(
gfile.GFile(model_path, 'rb').read())
self.assertCountEqual(op_set, ['ADD'])
