wehub-resource-sync
235 lines
8.4 KiB
Python
235 lines
8.4 KiB
Python
# Copyright 2019 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.
|
|
# ==============================================================================
|
|
"""Converts a model's graph def into a tflite model with MLIR-based conversion."""
|
|
import os
|
|
import shlex
|
|
import subprocess
|
|
import tempfile
|
|
|
|
import numpy as np
|
|
|
|
from tensorflow.lite.python import lite
|
|
from tensorflow.lite.python import test_util as tflite_test_util
|
|
from tensorflow.lite.testing import zip_test_utils
|
|
from tensorflow.python.platform import resource_loader
|
|
from tensorflow.python.saved_model import signature_constants
|
|
|
|
|
|
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", {})
|
|
extra_convert_options = kwargs.get("extra_convert_options",
|
|
zip_test_utils.ExtraConvertOptions())
|
|
tflite_model = None
|
|
log = ""
|
|
|
|
signature_key = signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY
|
|
converter = lite.TFLiteConverterV2.from_saved_model(
|
|
saved_model_dir, [signature_key])
|
|
converter.allow_custom_ops = extra_convert_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(
|
|
[lite.OpsSet.TFLITE_BUILTINS, lite.OpsSet.SELECT_TF_OPS])
|
|
|
|
if options.enable_dynamic_update_slice:
|
|
converter._experimental_enable_dynamic_update_slice = True # pylint: disable=protected-access
|
|
|
|
converter.unfold_batchmatmul = options.unfold_batchmatmul
|
|
|
|
if test_params.get("dynamic_range_quantize", False):
|
|
converter.optimizations = [lite.Optimize.DEFAULT]
|
|
|
|
if options.experimental_low_bit_qat:
|
|
converter._experimental_low_bit_qat = ( # pylint: disable=protected-access
|
|
True
|
|
)
|
|
|
|
if options.experimental_unsafe_single_batch_rank_reduction:
|
|
converter._experimental_unsafe_single_batch_rank_reduction = ( # pylint: disable=protected-access
|
|
True
|
|
)
|
|
|
|
if test_params.get("fully_quantize", False):
|
|
converter.optimizations = [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 = [
|
|
lite.OpsSet
|
|
.EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8
|
|
]
|
|
else:
|
|
converter.target_spec.supported_ops = [
|
|
lite.OpsSet.TFLITE_BUILTINS_INT8
|
|
]
|
|
|
|
converter.representative_dataset = representative_dataset_gen
|
|
if extra_convert_options.inference_input_type:
|
|
converter.inference_input_type = (
|
|
extra_convert_options.inference_input_type)
|
|
|
|
if extra_convert_options.inference_output_type:
|
|
converter.inference_output_type = (
|
|
extra_convert_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_file(graph_def_filename,
|
|
input_tensors,
|
|
output_tensors,
|
|
quantization_params=None,
|
|
additional_flags=""):
|
|
"""Convert a graphdef file into a tflite model with MLIR-based conversion.
|
|
|
|
NOTE: this currently shells out to the MLIR binary binary, but we would like
|
|
convert to Python API tooling in the future.
|
|
|
|
Args:
|
|
graph_def_filename: A GraphDef file.
|
|
input_tensors: List of input tensor tuples `(name, shape, type)`. name
|
|
should be a string. shape should be a tuple of integers. type should be a
|
|
string, for example 'DT_FLOAT'
|
|
output_tensors: List of output tensors (names).
|
|
quantization_params: parameters `(inference_type, min_values, max_values)`
|
|
to quantize the model.
|
|
additional_flags: A string of additional command line flags to be passed to
|
|
MLIR converter.
|
|
|
|
Returns:
|
|
output tflite model, log_txt from conversion
|
|
or None, log_txt if it did not convert properly.
|
|
"""
|
|
bin_path = resource_loader.get_path_to_datafile(
|
|
"../../compiler/mlir/lite/tf_tfl_translate"
|
|
)
|
|
if not os.path.exists(bin_path):
|
|
bin_path = resource_loader.get_path_to_datafile(
|
|
"../../../../compiler/mlir/lite/tf_tfl_translate"
|
|
)
|
|
|
|
with tempfile.NamedTemporaryFile() as output_file:
|
|
input_shapes = []
|
|
for input_tensor in input_tensors:
|
|
shape = input_tensor[1]
|
|
input_shapes.append(",".join([str(dim) for dim in shape]))
|
|
input_shapes_str = ":".join(input_shapes)
|
|
|
|
input_types = ",".join([x[2] for x in input_tensors])
|
|
|
|
cmd_list = [
|
|
bin_path,
|
|
"-tf-input-arrays=" + ",".join([x[0] for x in input_tensors]),
|
|
"-tf-input-data-types=" + input_types,
|
|
"-tf-input-shapes=" + input_shapes_str,
|
|
"-tf-output-arrays=" + ",".join(output_tensors),
|
|
]
|
|
if quantization_params is not None:
|
|
min_vals = ",".join([str(val) for val in quantization_params[1]])
|
|
max_vals = ",".join([str(val) for val in quantization_params[2]])
|
|
cmd_list.extend([
|
|
"-tf-inference-type=" + quantization_params[0],
|
|
"-tf-input-min-values=" + min_vals,
|
|
"-tf-input-max-values=" + max_vals,
|
|
"-emit-quant-adaptor-ops",
|
|
])
|
|
|
|
if additional_flags:
|
|
cmd_list.extend(shlex.split(additional_flags))
|
|
|
|
cmd_list.extend([graph_def_filename, "-o", output_file.name])
|
|
|
|
try:
|
|
with tempfile.NamedTemporaryFile() as stdout_file:
|
|
with tempfile.NamedTemporaryFile() as stderr_file:
|
|
result = subprocess.run(
|
|
cmd_list, stdout=stdout_file, stderr=stderr_file, check=False
|
|
)
|
|
exit_code = result.returncode
|
|
stdout_file.seek(0)
|
|
stderr_file.seek(0)
|
|
stdout = stdout_file.read().decode("utf-8", errors="replace")
|
|
stderr = stderr_file.read().decode("utf-8", errors="replace")
|
|
except FileNotFoundError:
|
|
exit_code = 127
|
|
stdout = ""
|
|
stderr = "Command not found: " + bin_path
|
|
except PermissionError:
|
|
exit_code = 126
|
|
stdout = ""
|
|
stderr = "Permission denied: " + bin_path
|
|
|
|
log = (
|
|
" ".join(cmd_list)
|
|
+ " exited with code %d" % exit_code
|
|
+ "\n------------------\n"
|
|
+ stdout
|
|
+ stderr
|
|
)
|
|
return (None if exit_code != 0 else output_file.read()), log
|