def toco_options(data_types,
input_arrays,
output_arrays,
shapes,
extra_toco_options=None):
"""Create TOCO options to process a model.
Args:
data_types: input and inference types used by TOCO.
input_arrays: names of the input tensors
output_arrays: name of the output tensors
shapes: shapes of the input tensors
extra_toco_options: additional toco options
Returns:
the options in a string.
"""
if extra_toco_options is None:
extra_toco_options = zip_test_utils.ExtraTocoOptions()
shape_str = ":".join([",".join(str(y) for y in x) for x in shapes if x])
inference_type = "FLOAT"
# TODO(ahentz): if we get multi-input quantization to work we need this
# to change
if data_types[0] == "QUANTIZED_UINT8":
inference_type = "QUANTIZED_UINT8"
s = (" --input_data_types=%s" % ",".join(data_types) +
" --inference_type=%s" % inference_type +
" --input_format=TENSORFLOW_GRAPHDEF" + " --output_format=TFLITE" +
" --input_arrays=%s" % ",".join(input_arrays) +
" --output_arrays=%s" % ",".join(output_arrays))
if shape_str:
s += (" --input_shapes=%s" % shape_str)
if extra_toco_options.drop_control_dependency:
s += " --drop_control_dependency"
if extra_toco_options.allow_custom_ops:
s += " --allow_custom_ops"
if extra_toco_options.rnn_states:
s += (" --rnn_states='" + extra_toco_options.rnn_states + "'")
if extra_toco_options.split_tflite_lstm_inputs is not None:
if extra_toco_options.split_tflite_lstm_inputs:
s += " --split_tflite_lstm_inputs=true"
else:
s += " --split_tflite_lstm_inputs=false"
return s
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 toco_convert(options, graph_def, input_tensors, output_tensors, **kwargs):
"""Convert a model's graph def into a tflite model.
NOTE: this currently shells out to the toco binary, but we would like
convert to Python API tooling in the future.
Args:
options: An 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 options to be passed.
Returns:
output tflite model, log_txt from conversion
or None, log_txt if it did not convert properly.
"""
# Convert ophint ops if presented.
graph_def = tf.compat.v1.lite.experimental.convert_op_hints_to_stubs(
graph_def=graph_def)
graph_def_str = graph_def.SerializeToString()
extra_toco_options = kwargs.get("extra_toco_options",
zip_test_utils.ExtraTocoOptions())
test_params = kwargs.get("test_params", {})
input_arrays = [x[0] for x in input_tensors]
data_types = [zip_test_utils.TF_TYPE_INFO[x[2]][1] for x in input_tensors]
if test_params.get("fully_quantize", False):
# Read the input range for the representative dataset from parameters.
min_value, max_value = test_params.get("input_range", (-1, 1))
with tempfile.NamedTemporaryFile() as graphdef_file:
graphdef_file.write(graph_def_str)
graphdef_file.flush()
input_shapes = zip_test_utils.get_input_shapes_map(input_tensors)
converter = tf.compat.v1.lite.TFLiteConverter.from_frozen_graph(
graphdef_file.name, input_arrays, output_tensors, input_shapes)
# TODO(b/145313371): Evaluate should we make it work with the new
# converter.
# Note: Currently this line is a non-functional change because the new
# converter is disabled by default. Since this code path doesn't work
# with new converter yet, it's explicitly disabled for easier testing.
converter.experimental_new_converter = False
def representative_dataset(input_tensors):
calibration_inputs = []
for _, shape, _ in input_tensors:
if shape:
dims = [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)
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)
else:
converter.inference_output_type = tf.int8
try:
tflite_model = converter.convert()
return tflite_model, ""
except Exception as e:
log = "{0}\n{1}".format(str(e), traceback.format_exc())
return None, log
else:
opts = toco_options(data_types=data_types,
input_arrays=input_arrays,
shapes=[x[1] for x in input_tensors],
output_arrays=output_tensors,
extra_toco_options=extra_toco_options)
with tempfile.NamedTemporaryFile() as graphdef_file, \
tempfile.NamedTemporaryFile() as output_file, \
tempfile.NamedTemporaryFile("w+") as stdout_file:
graphdef_file.write(graph_def_str)
graphdef_file.flush()
# TODO(aselle): Switch this to subprocess at some point.
if options.run_with_flex:
opts += " --enable_select_tf_ops --force_select_tf_ops"
cmd = ("%s --input_file=%s --output_file=%s %s > %s 2>&1" %
(options.toco, graphdef_file.name, output_file.name, opts,
stdout_file.name))
exit_code = os.system(cmd)
log = (cmd + "exited with code %d" % exit_code +
"\n------------------\n" + stdout_file.read())
return (None if exit_code != 0 else output_file.read()), log
def toco_convert(options, graph_def, input_tensors, output_tensors, **kwargs):
"""Convert a model's graph def into a tflite model.
NOTE: this currently shells out to the toco binary, but we would like
convert to Python API tooling in the future.
Args:
options: An 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 options to be passed.
Returns:
output tflite model, log_txt from conversion
or None, log_txt if it did not convert properly.
"""
# Convert ophint ops if presented.
graph_def = tf.compat.v1.lite.experimental.convert_op_hints_to_stubs(
graph_def=graph_def)
graph_def_str = graph_def.SerializeToString()
extra_toco_options = kwargs.get("extra_toco_options",
zip_test_utils.ExtraTocoOptions())
test_params = kwargs.get("test_params", {})
input_arrays = [x[0] for x in input_tensors]
data_types = [zip_test_utils.TF_TYPE_INFO[x[2]][1] for x in input_tensors]
fully_quantize = test_params.get("fully_quantize", False)
dynamic_range_quantize = test_params.get("dynamic_range_quantize", False)
if dynamic_range_quantize or fully_quantize:
with tempfile.NamedTemporaryFile() as graphdef_file:
graphdef_file.write(graph_def_str)
graphdef_file.flush()
input_shapes = zip_test_utils.get_input_shapes_map(input_tensors)
converter = tf.compat.v1.lite.TFLiteConverter.from_frozen_graph(
graphdef_file.name, input_arrays, output_tensors, input_shapes)
converter.experimental_new_quantizer = options.mlir_quantizer
converter.optimizations = [tf.lite.Optimize.DEFAULT]
if fully_quantize:
# 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 = [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)
else:
if test_params.get("quant_16x8", False):
converter.inference_output_type = tf.int16
else:
converter.inference_output_type = tf.int8
try:
tflite_model = converter.convert()
return tflite_model, ""
except Exception as e:
log = "{0}\n{1}".format(str(e), traceback.format_exc())
return None, log
else:
opts = toco_options(data_types=data_types,
input_arrays=input_arrays,
shapes=[x[1] for x in input_tensors],
output_arrays=output_tensors,
extra_toco_options=extra_toco_options)
with tempfile.NamedTemporaryFile() as graphdef_file, \
tempfile.NamedTemporaryFile() as output_file, \
tempfile.NamedTemporaryFile("w+") as stdout_file:
graphdef_file.write(graph_def_str)
graphdef_file.flush()
if options.run_with_flex:
opts += " --enable_select_tf_ops --force_select_tf_ops"
cmd = ("%s --input_file=%s --output_file=%s %s > %s 2>&1" %
(options.toco, graphdef_file.name, output_file.name, opts,
stdout_file.name))
exit_code = os.system(cmd)
log = (cmd + "exited with code %d" % exit_code +
"\n------------------\n" + stdout_file.read())
return (None if exit_code != 0 else output_file.read()), 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
