wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 3 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/lite/python/BUILD
+731
View File
@@ -0,0 +1,731 @@
load("@flatbuffers//:build_defs.bzl", "flatbuffer_py_library")
load("@rules_shell//shell:sh_test.bzl", "sh_test")
load("@xla//third_party/rules_python/python:py_binary.bzl", "py_binary")
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
load("//tensorflow:pytype.default.bzl", "pytype_strict_contrib_test", "pytype_strict_library")
load("//tensorflow:tensorflow.bzl", "if_oss", "py_test")
load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable", "pywrap_binaries", "pywrap_library")
load("//tensorflow/core/platform:build_config_root.bzl", "if_pywrap")
load("//tensorflow/lite:special_rules.bzl", "internal_visibility_allowlist")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
default_visibility = [
"//tensorflow:__subpackages__",
"//tensorflow:internal",
"//third_party/odml/infra/genai/conversion:__subpackages__",
"//third_party/odml/litert:__subpackages__",
"//third_party/odml/model_customization/quantization:__subpackages__",
"//third_party/py/ai_edge_torch:__subpackages__",
"//third_party/py/litert_torch:__subpackages__",
"//third_party/py/tensorflow_federated:__subpackages__",
"//third_party/tflite_micro:__subpackages__",
],
licenses = ["notice"],
)
exports_files([
"tflite_convert.py",
"pywrap_tflite_common.json",
"pywrap_tflite_common.lds",
"pywrap_tflite_common_darwin.lds",
])
flatbuffer_py_library(
name = "schema_py",
srcs = ["//tensorflow/compiler/mlir/lite/schema:schema.fbs"],
)
flatbuffer_py_library(
name = "conversion_metadata_schema_py",
srcs = ["//tensorflow/compiler/mlir/lite/schema:conversion_metadata.fbs"],
)
py_library(
name = "interpreter",
srcs = [
"interpreter.py",
],
compatible_with = get_compatible_with_portable(),
strict_deps = True,
visibility = ["//visibility:public"],
deps = [
"//tensorflow/lite/python/interpreter_wrapper:_pywrap_tensorflow_interpreter_wrapper",
"//tensorflow/lite/python/metrics",
"//tensorflow/python/util:tf_export",
"//third_party/py/numpy",
],
)
py_test(
name = "interpreter_test",
srcs = ["interpreter_test.py"],
data = [
"//tensorflow/lite:testdata/sparse_tensor.bin",
"//tensorflow/lite/python/testdata:interpreter_test_data",
"//tensorflow/lite/python/testdata:test_delegate.so",
],
# Static linking is required because this loads a cc_binary as a shared
# library, which would otherwise create ODR violations.
# copybara:uncomment linking_mode = "static",
strict_deps = True,
tags = [
"no_oss", # TODO(b/190842754): Enable test in OSS.
],
deps = [
":interpreter",
#internal proto upb dep
"//third_party/py/numpy",
"//tensorflow:tensorflow_py",
"//tensorflow/lite/python:lite",
"//tensorflow/lite/python/metrics",
"//tensorflow/lite/python/testdata:_pywrap_test_registerer",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/platform:resource_loader",
],
)
py_binary(
name = "tflite_convert",
srcs = ["tflite_convert.py"],
strict_deps = True,
visibility = ["//visibility:public"],
deps = [
":tflite_convert_main_lib",
"//tensorflow:tensorflow_py",
"//tensorflow/lite/python:lite",
"@absl_py//absl:app",
],
)
py_library(
name = "tflite_convert_main_lib",
srcs = ["tflite_convert.py"],
strict_deps = True,
visibility = ["//visibility:public"],
deps = [
":convert",
"//tensorflow:tensorflow_py",
"//tensorflow/lite/toco:toco_flags_proto_py",
"//tensorflow/lite/toco/logging:gen_html",
"//tensorflow/python:tf2",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/platform:gfile",
"//tensorflow/python/util:keras_deps",
"@absl_py//absl:app",
],
)
py_library(
name = "test_util",
srcs = ["test_util.py"],
strict_deps = True,
deps = [
":lite",
":schema_py",
":schema_util",
"//tensorflow/lite/tools:visualize_lib",
],
)
py_test(
name = "test_util_test",
srcs = ["test_util_test.py"],
data = [
"//tensorflow/lite:testdata/add.bin",
"//tensorflow/lite:testdata/softplus_flex.bin",
],
strict_deps = True,
deps = [
":test_util",
#internal proto upb dep
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/platform:gfile",
"//tensorflow/python/platform:resource_loader",
],
)
py_test(
name = "tflite_convert_test",
srcs = ["tflite_convert_test.py"],
data = [
":tflite_convert.par",
"@tflite_mobilenet_ssd_quant_protobuf//:tflite_graph.pb",
],
# Increased thread count for reducing timeout failures.
shard_count = 10,
strict_deps = True,
tags = [
"no_oss",
"no_windows",
],
deps = [
":convert",
":test_util",
":tflite_convert_main_lib",
#internal proto upb dep
"//third_party/py/numpy",
"//tensorflow:tensorflow_py",
"//tensorflow/core:protos_all_py",
"//tensorflow/python:tf2",
"//tensorflow/python/client:session",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:importer",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:random_ops",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/platform:gfile",
"//tensorflow/python/platform:resource_loader",
"//tensorflow/python/saved_model",
"//tensorflow/python/saved_model:save",
"//tensorflow/python/trackable:autotrackable",
"//tensorflow/python/training:training_util",
"@absl_py//absl/testing:parameterized",
],
)
py_library(
name = "lite",
srcs = ["lite.py"],
strict_deps = True,
tags = [
"ignore_for_dep=third_party.py.keras",
],
visibility = ["//visibility:public"],
deps = [
":conversion_metadata_schema_py",
":convert",
":convert_phase",
":convert_saved_model",
":interpreter",
":lite_constants",
":op_hint",
":util",
"//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_py",
"//tensorflow/compiler/mlir/quantization/tensorflow/python:representative_dataset",
"//tensorflow/core:protos_all_py",
"//tensorflow/lite/experimental/microfrontend:audio_microfrontend_py",
"//tensorflow/lite/profiling/proto:model_runtime_info_py",
"//tensorflow/lite/profiling/proto:profiling_info_py",
"//tensorflow/lite/python/metrics",
"//tensorflow/lite/python/optimize:calibrator",
"//tensorflow/lite/tools:flatbuffer_utils",
"//tensorflow/lite/tools/optimize/debugging/python:debugger",
"//tensorflow/python/client:session",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/eager:function",
"//tensorflow/python/framework:byte_swap_tensor",
"//tensorflow/python/framework:convert_to_constants",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:importer",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_spec",
"//tensorflow/python/framework:versions",
"//tensorflow/python/platform:gfile",
"//tensorflow/python/saved_model:load",
"//tensorflow/python/saved_model:loader",
"//tensorflow/python/saved_model:save",
"//tensorflow/python/saved_model:save_options",
"//tensorflow/python/saved_model:signature_constants",
"//tensorflow/python/saved_model:tag_constants",
"//tensorflow/python/util:deprecation",
"//tensorflow/python/util:keras_deps",
"//tensorflow/python/util:nest",
"//tensorflow/python/util:tf_export",
"@absl_py//absl/logging",
],
)
py_test(
name = "lite_test",
srcs = ["lite_test.py"],
data = [
"//tensorflow/lite/python/testdata:control_flow_v1.pbtxt",
"@tflite_mobilenet_ssd_quant_protobuf//:tflite_graph.pb",
],
shard_count = 4,
strict_deps = True,
tags = [
"no_windows",
],
deps = [
":conversion_metadata_schema_py",
":convert",
":interpreter",
":lite",
":lite_constants",
":schema_py",
":util",
#internal proto upb dep
"//third_party/py/numpy",
"//tensorflow:tensorflow_py",
"//tensorflow/python/client:session",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:convert_to_constants",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/framework:versions",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:logging_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:nn_ops",
"//tensorflow/python/ops:random_ops",
"//tensorflow/python/ops:variable_scope",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/platform:gfile",
"//tensorflow/python/platform:resource_loader",
"//tensorflow/python/saved_model",
"//tensorflow/python/training:training_util",
"@absl_py//absl/testing:parameterized",
],
)
py_test(
name = "lite_v2_test",
srcs = ["lite_v2_test.py"],
data = [
"//tensorflow/lite/python/testdata:test_delegate.so",
"//tensorflow/lite/python/testdata/control_flow_v1_saved_model:saved_model.pb",
],
exec_properties = if_oss(
None,
select({
"@bazel_tools//tools/cpp:asan_build": {"cpp_link.mem": "20g"},
"//conditions:default": None,
}),
),
shard_count = 18,
strict_deps = True,
tags = [
"no_windows",
],
deps = [
":conversion_metadata_schema_py",
":convert",
":interpreter",
":lite",
":lite_v2_test_util",
":schema_py",
":test_util",
":util",
#internal proto upb dep
"//third_party/py/numpy",
"//tensorflow:tensorflow_py",
"//tensorflow/compiler/mlir/quantization/stablehlo:quantization_options_proto_py",
"//tensorflow/lite/python/testdata:_pywrap_test_registerer",
"//tensorflow/lite/python/testdata:double_op",
"//tensorflow/lite/tools:flatbuffer_utils",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/framework:versions",
"//tensorflow/python/lib/io:file_io",
"//tensorflow/python/ops:map_ops",
"//tensorflow/python/ops:rnn",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/platform:resource_loader",
"//tensorflow/python/saved_model",
"//tensorflow/python/saved_model:loader",
"//tensorflow/python/saved_model:save",
"//tensorflow/python/saved_model:save_options",
"//tensorflow/python/trackable:autotrackable",
"@absl_py//absl/testing:parameterized",
"@pypi//jax",
],
)
py_library(
name = "lite_v2_test_util",
testonly = 1,
srcs = ["lite_v2_test_util.py"],
strict_deps = True,
tags = [
"no_windows",
],
deps = [
":interpreter",
":lite",
"//tensorflow:tensorflow_py_no_contrib",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:tensor_spec",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/trackable:autotrackable",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
)
py_test(
name = "lite_flex_test",
srcs = ["lite_flex_test.py"],
strict_deps = True,
deps = [
":convert",
":interpreter",
":lite",
":test_util",
#internal proto upb dep
"//third_party/py/numpy",
"//tensorflow:tensorflow_py",
"//tensorflow/core:protos_all_py",
"//tensorflow/lite/python/testdata:double_op",
"//tensorflow/python/client:session",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:importer",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:list_ops",
"//tensorflow/python/ops:nn_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/saved_model",
"//tensorflow/python/trackable:autotrackable",
"@absl_py//absl/testing:parameterized",
],
)
py_library(
name = "util",
srcs = ["util.py"],
strict_deps = True,
visibility = internal_visibility_allowlist(),
deps = [
":conversion_metadata_schema_py",
":op_hint",
":schema_py",
":schema_util",
":tflite_keras_util",
"//tensorflow/core:protos_all_py",
"//tensorflow/lite/tools:flatbuffer_utils",
"//tensorflow/python/eager:function",
"//tensorflow/python/framework:convert_to_constants",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:error_interpolation",
"//tensorflow/python/grappler:tf_optimizer",
"//tensorflow/python/training:saver",
"//third_party/py/numpy",
"@absl_py//absl/logging",
"@pypi//flatbuffers",
],
)
py_test(
name = "util_test",
srcs = ["util_test.py"],
strict_deps = True,
tags = [
"no_windows",
],
deps = [
":util",
#internal proto upb dep
"//third_party/py/numpy",
"//tensorflow:tensorflow_py",
"//tensorflow/lite/python:lite",
"//tensorflow/lite/tools:flatbuffer_utils",
"//tensorflow/python/client:session",
"//tensorflow/python/framework:convert_to_constants",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:while_loop",
"//tensorflow/python/platform:client_testlib",
"@absl_py//absl/testing:parameterized",
],
)
py_library(
name = "tflite_keras_util",
srcs = [
"tflite_keras_util.py",
],
strict_deps = True,
deps = [
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:tensor_spec",
"//tensorflow/python/util:compat",
"//tensorflow/python/util:keras_deps",
"//tensorflow/python/util:nest",
],
)
py_library(
name = "lite_constants",
srcs = ["lite_constants.py"],
strict_deps = True,
deps = [
"//tensorflow/compiler/mlir/lite:converter_flags_proto_py",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/util:all_util",
"//tensorflow/python/util:tf_export",
],
)
pytype_strict_library(
name = "convert",
srcs = ["convert.py"],
visibility = ["//visibility:public"],
deps = [
":convert_phase",
":lite_constants",
":util",
"//tensorflow/compiler/mlir/lite:converter_flags_proto_py",
"//tensorflow/compiler/mlir/lite:model_flags_proto_py",
"//tensorflow/compiler/mlir/lite:types_proto_py",
"//tensorflow/compiler/mlir/lite/metrics:converter_error_data_proto_py",
"//tensorflow/compiler/mlir/lite/python:wrap_converter",
"//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_py",
"//tensorflow/compiler/mlir/quantization/stablehlo:quantization_options_proto_py",
"//tensorflow/lite/python/metrics:metrics_wrapper",
"//tensorflow/lite/tools:flatbuffer_utils",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/util:deprecation",
"//tensorflow/python/util:tf_export",
] + if_pywrap(
if_true = [":pywrap_tflite"],
),
)
py_library(
name = "op_hint",
srcs = ["op_hint.py"],
strict_deps = True,
visibility = ["//visibility:public"],
deps = [
"//tensorflow/core:protos_all_py",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:graph_util",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_util",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/util:all_util",
"//tensorflow/python/util:compat",
"//tensorflow/python/util:deprecation",
"//tensorflow/python/util:tf_export",
],
)
pytype_strict_contrib_test(
name = "convert_test",
srcs = ["convert_test.py"],
deps = [
":convert",
":interpreter",
":op_hint",
"//tensorflow/compiler/mlir/lite:converter_flags_proto_py",
"//tensorflow/compiler/mlir/lite/metrics:converter_error_data_proto_py",
"//tensorflow/lite/python/metrics:metrics_wrapper",
"//tensorflow/python/client:session",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:graph_util",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_stack",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
] + if_pywrap(
if_true = [":pywrap_tflite"],
),
)
py_library(
name = "convert_saved_model",
srcs = ["convert_saved_model.py"],
strict_deps = True,
visibility = [
"//tensorflow/lite:__subpackages__",
],
deps = [
":convert_phase",
":util",
"//tensorflow/python/client:session",
"//tensorflow/python/framework:graph_util",
"//tensorflow/python/framework:ops",
"//tensorflow/python/platform:tf_logging",
"//tensorflow/python/saved_model",
"//tensorflow/python/saved_model:constants",
"//tensorflow/python/saved_model:loader",
],
)
py_test(
name = "convert_saved_model_test",
srcs = ["convert_saved_model_test.py"],
strict_deps = True,
tags = [
"no_windows",
],
visibility = ["//visibility:public"],
deps = [
":convert_saved_model",
#internal proto upb dep
"//tensorflow/python/client:session",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/layers",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops/losses",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/platform:test",
"//tensorflow/python/saved_model",
"//tensorflow/python/saved_model:signature_constants",
"//tensorflow/python/saved_model:tag_constants",
],
)
py_binary(
name = "convert_file_to_c_source",
srcs = ["convert_file_to_c_source.py"],
strict_deps = True,
visibility = ["//visibility:public"],
deps = [
":util",
"@absl_py//absl:app",
"@absl_py//absl/flags",
],
)
sh_test(
name = "convert_file_to_c_source_test",
srcs = ["convert_file_to_c_source_test.sh"],
data = [":convert_file_to_c_source"],
)
py_library(
name = "schema_util",
srcs = ["schema_util.py"],
strict_deps = True,
visibility = ["//tensorflow/lite/schema:utils_friends"],
deps = [
"//tensorflow/python/util:all_util",
],
)
# Use py_library since the metrics module is imported in a try-except block,
# which doesn't work with the pytype_strict_library.
py_library(
name = "convert_phase",
srcs = ["convert_phase.py"],
strict_deps = True,
visibility = ["//tensorflow/lite:__subpackages__"],
deps = [
"//tensorflow/compiler/mlir/lite/metrics:converter_error_data_proto_py",
"//tensorflow/lite/python/metrics",
],
)
py_library(
name = "analyzer",
srcs = [
"analyzer.py",
],
strict_deps = True,
visibility = ["//visibility:public"],
deps = [
"//tensorflow/compiler/mlir/lite/python:wrap_converter",
"//tensorflow/python/util:tf_export",
] + if_pywrap(
if_false = [
"//tensorflow/lite/python/analyzer_wrapper:_pywrap_analyzer_wrapper",
],
if_true = [
":pywrap_tflite",
],
),
)
py_test(
name = "analyzer_test",
srcs = ["analyzer_test.py"],
data = [
"//tensorflow/lite:testdata/add.bin",
"//tensorflow/lite:testdata/conv_huge_im2col.bin",
"//tensorflow/lite:testdata/multi_add_flex.bin",
],
strict_deps = True,
deps = [
":analyzer",
#internal proto upb dep
"//tensorflow:tensorflow_py",
"//tensorflow/lite/python:lite",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/platform:resource_loader",
"//tensorflow/python/trackable:autotrackable",
],
)
# Use pywrap_library to avoid duplicate registration of pybind11 modules.
# A great example on how to use pywrap_library is
# https://github.com/vam-google/symbol-locations/blob/main/pybind/BUILD
# The following pywrap_library is used by LiteRT repo to avoid shared links provided
# by Tensorflow under tensorflow/python:_pywrap_tensorflow
# This isolate LiteRT's pybind11 dependencies. To use, add pybind deps under pywrap_tflite
# and refer pywrap_tflite to any target that needsd to selected shared objects.
py_library(
name = "tflite_pywrap_deps",
strict_deps = True,
visibility = ["//visibility:public"],
deps = [
"//tensorflow/lite/experimental/genai:pywrap_genai_ops",
"//tensorflow/lite/python/analyzer_wrapper:_pywrap_analyzer_wrapper",
"//tensorflow/lite/python/interpreter_wrapper:_pywrap_tensorflow_interpreter_wrapper",
"//tensorflow/lite/python/metrics:_pywrap_tensorflow_lite_metrics_wrapper",
"//tensorflow/lite/python/optimize:_pywrap_tensorflow_lite_calibration_wrapper",
"//tensorflow/lite/testing:_pywrap_string_util",
"//tensorflow/lite/tools/optimize/python:_pywrap_modify_model_interface",
"//tensorflow/lite/tools/optimize/sparsity:format_converter_wrapper_pybind11",
],
)
pywrap_library(
name = "pywrap_tflite",
common_lib_def_files_or_filters = {
"tensorflow/lite/python/pywrap_tflite_common": "pywrap_tflite_common.json",
},
common_lib_version_scripts = {
"tensorflow/lite/python/pywrap_tflite_common": select({
"@bazel_tools//src/conditions:windows": None,
"@bazel_tools//src/conditions:darwin": "pywrap_tflite_common_darwin.lds",
"//conditions:default": "pywrap_tflite_common.lds",
}),
},
pywrap_count = 8,
visibility = ["//visibility:public"],
deps = [
":tflite_pywrap_deps",
],
)
pywrap_binaries(
name = "pywrap_tflite_binaries",
dep = ":pywrap_tflite",
)
tensorflow/lite/python/analyzer.py
+107
View File
@@ -0,0 +1,107 @@
# Copyright 2021 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.
# ==============================================================================
"""This tool analyzes a TensorFlow Lite graph."""
import os
# pylint: disable=g-import-not-at-top
if not os.path.splitext(__file__)[0].endswith(
os.path.join("tflite_runtime", "analyzer")):
# This file is part of tensorflow package.
from tensorflow.compiler.mlir.lite.python import wrap_converter
from tensorflow.lite.python.analyzer_wrapper import _pywrap_analyzer_wrapper as _analyzer_wrapper
from tensorflow.python.util.tf_export import tf_export as _tf_export
else:
# This file is part of tflite_runtime package.
from tflite_runtime import _pywrap_analyzer_wrapper as _analyzer_wrapper
def _tf_export(*x, **kwargs):
del x, kwargs
return lambda x: x
@_tf_export("lite.experimental.Analyzer")
class ModelAnalyzer():
"""Provides a collection of TFLite model analyzer tools.
Example:
```python
model = tf.keras.applications.MobileNetV3Large()
fb_model = tf.lite.TFLiteConverterV2.from_keras_model(model).convert()
tf.lite.experimental.Analyzer.analyze(model_content=fb_model)
# === TFLite ModelAnalyzer ===
#
# Your TFLite model has ‘1’ subgraph(s). In the subgraph description below,
# T# represents the Tensor numbers. For example, in Subgraph#0, the MUL op
# takes tensor #0 and tensor #19 as input and produces tensor #136 as output.
#
# Subgraph#0 main(T#0) -> [T#263]
# Op#0 MUL(T#0, T#19) -> [T#136]
# Op#1 ADD(T#136, T#18) -> [T#137]
# Op#2 CONV_2D(T#137, T#44, T#93) -> [T#138]
# Op#3 HARD_SWISH(T#138) -> [T#139]
# Op#4 DEPTHWISE_CONV_2D(T#139, T#94, T#24) -> [T#140]
# ...
```
WARNING: Experimental interface, subject to change.
"""
@staticmethod
def analyze(model_path=None,
model_content=None,
gpu_compatibility=False,
**kwargs):
"""Analyzes the given tflite_model with dumping model structure.
This tool provides a way to understand users' TFLite flatbuffer model by
dumping internal graph structure. It also provides additional features
like checking GPU delegate compatibility.
WARNING: Experimental interface, subject to change.
The output format is not guaranteed to stay stable, so don't
write scripts to this.
Args:
model_path: TFLite flatbuffer model path.
model_content: TFLite flatbuffer model object.
gpu_compatibility: Whether to check GPU delegate compatibility.
**kwargs: Experimental keyword arguments to analyze API.
Returns:
Print analyzed report via console output.
"""
if not model_path and not model_content:
raise ValueError("neither `model_path` nor `model_content` is provided")
if model_path:
print(f"=== {model_path} ===\n")
tflite_model = model_path
input_is_filepath = True
else:
print("=== TFLite ModelAnalyzer ===\n")
tflite_model = model_content
input_is_filepath = False
if kwargs.get("experimental_use_mlir", False):
print(
wrap_converter.wrapped_flat_buffer_file_to_mlir(
tflite_model, input_is_filepath
)
)
else:
print(
_analyzer_wrapper.ModelAnalyzer(tflite_model, input_is_filepath,
gpu_compatibility))
tensorflow/lite/python/analyzer_test.py
+209
View File
@@ -0,0 +1,209 @@
# Copyright 2021 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.
# ==============================================================================
"""Tests for analyzer package."""
import io
import sys
import tempfile
import tensorflow as tf
from tensorflow.lite.python import analyzer
from tensorflow.lite.python import lite
from tensorflow.python.framework import test_util
from tensorflow.python.platform import resource_loader
from tensorflow.python.platform import test
from tensorflow.python.trackable import autotrackable
class AnalyzerTest(test_util.TensorFlowTestCase):
def testTxt(self):
model_path = resource_loader.get_path_to_datafile('../testdata/add.bin')
mock_stdout = io.StringIO()
with test.mock.patch.object(sys, 'stdout', mock_stdout):
analyzer.ModelAnalyzer.analyze(model_path=model_path)
txt = mock_stdout.getvalue()
self.assertIn('Subgraph#0(T#1) -> [T#2]', txt)
self.assertIn('Op#0 ADD(T#1, T#1) -> [T#0]', txt)
self.assertIn('Op#1 ADD(T#0, T#1) -> [T#2]', txt)
self.assertNotIn('Your model looks compatible with GPU delegate', txt)
def testMlir(self):
model_path = resource_loader.get_path_to_datafile('../testdata/add.bin')
mock_stdout = io.StringIO()
with test.mock.patch.object(sys, 'stdout', mock_stdout):
analyzer.ModelAnalyzer.analyze(
model_path=model_path, experimental_use_mlir=True)
mlir = mock_stdout.getvalue()
self.assertIn(
'func @main(%arg0: tensor<1x8x8x3xf32> '
'{tf_saved_model.index_path = ["a"]}) -> '
'(tensor<1x8x8x3xf32> {tf_saved_model.index_path = ["x"]}) attributes '
'{tf.entry_function = {inputs = "input", outputs = "output"}, '
'tf_saved_model.exported_names = ["serving_default"]}', mlir)
self.assertIn(
'%0 = tfl.add %arg0, %arg0 {fused_activation_function = "NONE"} : '
'tensor<1x8x8x3xf32>', mlir)
self.assertIn(
'%1 = tfl.add %0, %arg0 {fused_activation_function = "NONE"} : '
'tensor<1x8x8x3xf32>', mlir)
self.assertIn('return %1 : tensor<1x8x8x3xf32>', mlir)
def testMlirHugeConst(self):
model_path = resource_loader.get_path_to_datafile(
'../testdata/conv_huge_im2col.bin')
mock_stdout = io.StringIO()
with test.mock.patch.object(sys, 'stdout', mock_stdout):
analyzer.ModelAnalyzer.analyze(
model_path=model_path, experimental_use_mlir=True)
mlir = mock_stdout.getvalue()
self.assertIn(
'%1 = "tfl.pseudo_const"() <{value = dense_resource<__elided__> : '
'tensor<3x3x3x8xf32>}> : () -> tensor<3x3x3x8xf32>', mlir)
def testTxtWithFlatBufferModel(self):
@tf.function(
input_signature=[tf.TensorSpec(shape=[None], dtype=tf.float32)])
def func(x):
return x + tf.cos(x)
converter = lite.TFLiteConverterV2.from_concrete_functions(
[func.get_concrete_function()], func)
fb_model = converter.convert()
mock_stdout = io.StringIO()
with test.mock.patch.object(sys, 'stdout', mock_stdout):
analyzer.ModelAnalyzer.analyze(model_content=fb_model)
txt = mock_stdout.getvalue()
self.assertIn('Subgraph#0 main(T#0) -> [T#2]', txt)
self.assertIn('Op#0 COS(T#0) -> [T#1]', txt)
self.assertIn('Op#1 ADD(T#0, T#1) -> [T#2]', txt)
def testMlirWithFlatBufferModel(self):
@tf.function(
input_signature=[tf.TensorSpec(shape=[None], dtype=tf.float32)])
def func(x):
return x + tf.cos(x)
converter = lite.TFLiteConverterV2.from_concrete_functions(
[func.get_concrete_function()], func)
fb_model = converter.convert()
mock_stdout = io.StringIO()
with test.mock.patch.object(sys, 'stdout', mock_stdout):
analyzer.ModelAnalyzer.analyze(
model_content=fb_model, experimental_use_mlir=True)
mlir = mock_stdout.getvalue()
self.assertIn('func @main(%arg0: tensor<?xf32>) -> tensor<?xf32>', mlir)
self.assertIn('%0 = "tfl.cos"(%arg0) : (tensor<?xf32>) -> tensor<?xf32>',
mlir)
self.assertIn(
'%1 = tfl.add %arg0, %0 {fused_activation_function = "NONE"} : '
'tensor<?xf32>', mlir)
self.assertIn('return %1 : tensor<?xf32', mlir)
def testTxtSignatureDefs(self):
with tempfile.TemporaryDirectory() as tmp_dir:
@tf.function(input_signature=[
tf.TensorSpec(shape=None, dtype=tf.float32),
tf.TensorSpec(shape=None, dtype=tf.float32)
])
def add(a, b):
return {'add_result': tf.add(a, b)}
@tf.function(input_signature=[
tf.TensorSpec(shape=None, dtype=tf.float32),
tf.TensorSpec(shape=None, dtype=tf.float32)
])
def sub(x, y):
return {'sub_result': tf.subtract(x, y)}
root = autotrackable.AutoTrackable()
root.f1 = add.get_concrete_function()
root.f2 = sub.get_concrete_function()
tf.saved_model.save(
root, tmp_dir, signatures={
'add': root.f1,
'sub': root.f2
})
converter = lite.TFLiteConverterV2.from_saved_model(tmp_dir)
fb_model = converter.convert()
mock_stdout = io.StringIO()
with test.mock.patch.object(sys, 'stdout', mock_stdout):
analyzer.ModelAnalyzer.analyze(model_content=fb_model)
txt = mock_stdout.getvalue()
self.assertIn("Your TFLite model has '2' signature_def(s).", txt)
self.assertIn("Signature#0 key: 'add'", txt)
self.assertIn(" 'a' : T#1", txt)
self.assertIn(" 'b' : T#0", txt)
self.assertIn(" 'add_result' : T#2", txt)
self.assertIn("Signature#1 key: 'sub'", txt)
self.assertIn(" 'x' : T#1_1", txt)
self.assertIn(" 'y' : T#1_0", txt)
self.assertIn(" 'sub_result' : T#1_2", txt)
def testTxtWithoutInput(self):
@tf.function()
def func():
return tf.cos(1.0)
converter = lite.TFLiteConverterV2.from_concrete_functions(
[func.get_concrete_function()], func)
fb_model = converter.convert()
mock_stdout = io.StringIO()
with test.mock.patch.object(sys, 'stdout', mock_stdout):
analyzer.ModelAnalyzer.analyze(model_content=fb_model)
txt = mock_stdout.getvalue()
self.assertIn('Subgraph#0 main() -> [T#0]', txt)
def testTxtWithEinsum(self):
@tf.function(input_signature=[
tf.TensorSpec(shape=[1, 100, 512], dtype=tf.float32),
tf.TensorSpec(shape=[512, 8, 64], dtype=tf.float32)
])
def func(lhs, rhs):
return tf.einsum('ABD,DNH->ABNH', lhs, rhs)
converter = lite.TFLiteConverterV2.from_concrete_functions(
[func.get_concrete_function()], func)
converter.unfold_batchmatmul = True
fb_model = converter.convert()
mock_stdout = io.StringIO()
with test.mock.patch.object(sys, 'stdout', mock_stdout):
analyzer.ModelAnalyzer.analyze(model_content=fb_model)
txt = mock_stdout.getvalue()
self.assertIn('Op#0 RESHAPE(T#1, T#4[512, 512]) -> [T#5]', txt)
self.assertIn('Op#1 TRANSPOSE(T#5, T#3[1, 0]) -> [T#6]', txt)
self.assertIn('Op#2 FULLY_CONNECTED(T#0, T#6, T#-1) -> [T#7]', txt)
self.assertIn('Op#3 RESHAPE(T#7, T#2[1, 100, 8, 64]) -> [T#8]', txt)
self.assertIn(
'T#2(einsum/Einsum) shape:[4], type:INT32 RO 16 bytes, '
'buffer: 3, data:[1, 100, 8, 64]', txt)
self.assertIn(
'T#3(einsum/Einsum2) shape:[2], type:INT32 RO 8 bytes, '
'buffer: 4, data:[1, 0]', txt)
self.assertIn(
'T#4(einsum/Einsum3) shape:[2], type:INT32 RO 8 bytes, '
'buffer: 5, data:[512, 512]', txt)
if __name__ == '__main__':
test.main()
tensorflow/lite/python/analyzer_wrapper/BUILD
+46
View File
@@ -0,0 +1,46 @@
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load("//tensorflow:tensorflow.default.bzl", "pybind_extension")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
default_visibility = ["//tensorflow:internal"],
licenses = ["notice"],
)
pybind_extension(
name = "_pywrap_analyzer_wrapper",
srcs = [
"analyzer_wrapper.cc",
],
common_lib_packages = [
"litert/python",
"tensorflow/lite/python",
],
enable_stub_generation = True,
pytype_srcs = [
"_pywrap_analyzer_wrapper.pyi",
],
wrap_py_init = True,
deps = [
":model_analyzer",
"@pybind11",
],
)
cc_library(
name = "model_analyzer",
srcs = ["model_analyzer.cc"],
hdrs = ["model_analyzer.h"],
visibility = ["//visibility:public"],
deps = [
"//tensorflow/compiler/mlir/lite/schema:schema_utils",
"//tensorflow/core:framework",
"//tensorflow/core/public:release_version",
"//tensorflow/lite:version",
"//tensorflow/lite/core:model_builder",
"//tensorflow/lite/core/api:error_reporter",
"//tensorflow/lite/schema:schema_fbs",
"@com_google_absl//absl/strings",
"@flatbuffers//:runtime_cc",
],
)
tensorflow/lite/python/analyzer_wrapper/_pywrap_analyzer_wrapper.pyi
+16
View File
@@ -0,0 +1,16 @@
# Copyright 2023 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.
# ==============================================================================
def ModelAnalyzer(arg0: str, arg1: bool, arg2: bool) -> str: ...
tensorflow/lite/python/analyzer_wrapper/analyzer_wrapper.cc
+32
View File
@@ -0,0 +1,32 @@
/* Copyright 2021 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.
==============================================================================*/
#include <string>
#include "pybind11/pybind11.h" // from @pybind11
#include "tensorflow/lite/python/analyzer_wrapper/model_analyzer.h"
PYBIND11_MODULE(_pywrap_analyzer_wrapper, m) {
m.def(
"ModelAnalyzer",
[](const std::string& model_path, bool input_is_filepath,
bool gpu_compatibility) {
return ::tflite::model_analyzer(model_path, input_is_filepath,
gpu_compatibility);
},
R"pbdoc(
Returns txt dump of the given TFLite file.
)pbdoc");
}
tensorflow/lite/python/analyzer_wrapper/model_analyzer.cc
+506
View File
@@ -0,0 +1,506 @@
/* Copyright 2021 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.
==============================================================================*/
#include <algorithm>
#include <cctype>
#include <cstdarg>
#include <cstddef>
#include <cstdint>
#include <cstdio>
#include <memory>
#include <sstream>
#include <string>
#include <vector>
#include "absl/strings/str_join.h"
#include "flatbuffers/vector.h" // from @flatbuffers
#include "tensorflow/compiler/mlir/lite/schema/schema_utils.h"
#include "tensorflow/core/public/release_version.h"
#include "tensorflow/lite/core/api/error_reporter.h"
#include "tensorflow/lite/core/model_builder.h"
#include "tensorflow/lite/schema/schema_generated.h"
#include "tensorflow/lite/version.h"
namespace tflite {
namespace {
const float kThreshold_zero_buffer_ratio = 10.0f;
constexpr char kSectionSplitter[] =
"---------------------------------------------------------------\n";
const int kMaxContentDumpCnt = 5;
// Returns string representation of the given tensor data up to 5 elements.
const std::string get_tensor_data_str(const tflite::Tensor* tensor,
const tflite::Model* model) {
std::stringstream ss;
auto buffer_idx = tensor->buffer();
if (buffer_idx != 0 && buffer_idx < model->buffers()->size()) {
auto* buffer = model->buffers()->Get(buffer_idx);
if (buffer->data() == nullptr) {
return "";
}
ss << "[";
if (buffer->data()->size() != 0) {
size_t type_size;
switch (tensor->type()) {
case tflite::TensorType_INT32:
case tflite::TensorType_UINT32:
case tflite::TensorType_FLOAT32:
type_size = 4;
break;
default:
type_size = 1;
}
int data_cnt = buffer->data()->size() / type_size;
for (int i = 0; i < std::min(kMaxContentDumpCnt, data_cnt); ++i) {
switch (tensor->type()) {
case tflite::TensorType_INT32: {
auto data =
reinterpret_cast<const int32_t*>(buffer->data()->data());
ss << data[i];
} break;
case tflite::TensorType_UINT32: {
auto data =
reinterpret_cast<const uint32_t*>(buffer->data()->data());
ss << data[i];
} break;
case tflite::TensorType_INT8: {
auto data = reinterpret_cast<const int8_t*>(buffer->data()->data());
ss << data[i];
} break;
case tflite::TensorType_UINT8: {
auto data =
reinterpret_cast<const uint8_t*>(buffer->data()->data());
ss << data[i];
} break;
case tflite::TensorType_FLOAT32: {
auto data = reinterpret_cast<const float*>(buffer->data()->data());
ss << data[i];
} break;
case tflite::TensorType_STRING: {
auto data = reinterpret_cast<const char*>(buffer->data()->data());
ss << data[i];
} break;
default:
ss << "??";
break;
}
if (i != data_cnt - 1) {
ss << ", ";
}
}
if (data_cnt > kMaxContentDumpCnt) {
ss << "...";
}
}
ss << "]";
}
return ss.str();
}
// Returns string representation of the given tensor of the subgraph.
const std::string tensor_str(const int tensor_idx, const int subgraph_idx,
const tflite::Model* model = nullptr) {
std::stringstream ss;
if (subgraph_idx != 0 && tensor_idx != -1)
ss << "T#" << subgraph_idx << "_" << tensor_idx;
else
ss << "T#" << tensor_idx;
if (model && tensor_idx != -1) {
const SubGraph* subgraph = model->subgraphs()->Get(subgraph_idx);
if (subgraph && subgraph->tensors()) {
auto tensor = subgraph->tensors()->Get(tensor_idx);
if (tensor && tensor->type() == tflite::TensorType_INT32) {
ss << get_tensor_data_str(tensor, model);
}
}
}
return ss.str();
}
// Returns string representation of the given subgraph.
const std::string subgraph_str(const int subgraph_idx) {
std::stringstream ss;
ss << "Subgraph#" << subgraph_idx;
return ss.str();
}
struct ModelStats {
// FlatBuffer buffer usage (in bytes) per subgraph.
std::vector<size_t> buffer_usage;
};
// Dump details of the given tensor.
void dump_tensor_detail(std::stringstream& out_stream,
const tflite::Tensor* tensor, const int tensor_idx,
const int subgraph_idx, const tflite::Model* model,
ModelStats* stats) {
out_stream << tensor_str(tensor_idx, subgraph_idx);
if (tensor->name()) {
out_stream << "(" << tensor->name()->str() << ") ";
}
// Prints `shape_signature` instead of `shape` if it's available since it
// supports dynamic shapes.
if (tensor->shape_signature()) {
out_stream << "shape_signature:[";
for (int i = 0; i < tensor->shape_signature()->size(); ++i) {
const int j = tensor->shape_signature()->Get(i);
out_stream << j;
if (i != tensor->shape_signature()->size() - 1) {
out_stream << ", ";
}
}
out_stream << "]";
} else if (tensor->shape()) {
out_stream << "shape:[";
for (int i = 0; i < tensor->shape()->size(); ++i) {
const int j = tensor->shape()->Get(i);
out_stream << j;
if (i != tensor->shape()->size() - 1) {
out_stream << ", ";
}
}
out_stream << "]";
} else {
out_stream << "shape:n/a";
}
out_stream << ", type:" << EnumNameTensorType(tensor->type());
// Dump buffer size of constant tensors.
auto buffer_idx = tensor->buffer();
if (buffer_idx != 0 && buffer_idx < model->buffers()->size()) {
auto* buffer = model->buffers()->Get(buffer_idx);
if (buffer->data() && buffer->data()->size() != 0) {
out_stream << " RO " << buffer->data()->size() << " bytes";
out_stream << ", buffer: " << buffer_idx;
out_stream << ", data:" << get_tensor_data_str(tensor, model);
stats->buffer_usage[subgraph_idx] += buffer->data()->size();
}
}
out_stream << "\n";
}
// Dump list of input or output tensors.
void dump_tensor_list(std::stringstream& out_stream,
const flatbuffers::Vector<int32_t>* tensors,
const int subgraph_idx,
const tflite::Model* model = nullptr,
bool verbose = false) {
if (tensors == nullptr) {
return;
}
for (int i = 0; i < tensors->size(); ++i) {
const int tensor_idx = tensors->Get(i);
if (verbose) {
out_stream << "tensor #" << tensor_idx;
} else {
out_stream << tensor_str(tensor_idx, subgraph_idx, model);
}
if (i != tensors->size() - 1) {
if (verbose) {
out_stream << " and ";
} else {
out_stream << ", ";
}
}
}
}
// Returns the string representation of the given OperatorCode.
const std::string get_op_name(const OperatorCode* op_code) {
auto builtin_code = GetBuiltinCode(op_code);
if (builtin_code != BuiltinOperator_CUSTOM) {
return EnumNameBuiltinOperator(builtin_code);
} else {
return op_code->custom_code()->str();
}
}
// Dump the given Operator node.
void dump_node(std::stringstream& out_stream, const int node_no,
const OperatorCode* op_code, const Operator* op,
const int subgraph_index, const ::tflite::Model* model) {
out_stream << "Op#" << node_no << " " << get_op_name(op_code);
out_stream << "(";
dump_tensor_list(out_stream, op->inputs(), subgraph_index, model);
if (GetBuiltinCode(op_code) == BuiltinOperator_CALL_ONCE) {
out_stream << subgraph_str(
op->builtin_options_as_CallOnceOptions()->init_subgraph_index());
} else if (GetBuiltinCode(op_code) == BuiltinOperator_IF) {
out_stream << ", Then: "
<< subgraph_str(op->builtin_options_as_IfOptions()
->then_subgraph_index());
out_stream << ", Else: "
<< subgraph_str(op->builtin_options_as_IfOptions()
->else_subgraph_index());
} else if (GetBuiltinCode(op_code) == BuiltinOperator_WHILE) {
out_stream << ", Cond: "
<< subgraph_str(op->builtin_options_as_WhileOptions()
->cond_subgraph_index());
out_stream << ", Body: "
<< subgraph_str(op->builtin_options_as_WhileOptions()
->body_subgraph_index());
}
out_stream << ") -> [";
dump_tensor_list(out_stream, op->outputs(), subgraph_index);
out_stream << "]\n";
}
// Dump the summary of the given TFLite flatbuffer model. It's printed at the
// beginning of the analyzer output.
void dump_model_summary(std::stringstream& out_stream,
const ::tflite::Model* model) {
auto* subgraphs = model->subgraphs();
out_stream
<< "Your TFLite model has '" << subgraphs->size()
<< "' subgraph(s). In the subgraph description below,\nT# represents the "
"Tensor numbers. ";
if (subgraphs->size() > 0 && subgraphs->Get(0)->operators()->size() > 0) {
const Operator* first_op = subgraphs->Get(0)->operators()->Get(0);
const OperatorCode* first_op_code =
model->operator_codes()->Get(first_op->opcode_index());
out_stream << "For example, in " << subgraph_str(0) << ", the "
<< get_op_name(first_op_code) << " op takes\n";
dump_tensor_list(out_stream, first_op->inputs(), 0, nullptr,
/*verbose=*/true);
out_stream << " as input and produces ";
dump_tensor_list(out_stream, first_op->outputs(), 0, nullptr,
/*verbose=*/true);
out_stream << " as output.\n\n";
}
}
// Dump the signature definitions of the given TFLite flatbuffer model.
void dump_model_signature_defs(std::stringstream& out_stream,
const ::tflite::Model* model) {
auto* signatures = model->signature_defs();
if (signatures == nullptr || signatures->size() == 0) {
return;
}
out_stream << kSectionSplitter;
out_stream << "Your TFLite model has '" << signatures->size()
<< "' signature_def(s).\n\n";
for (int i = 0; i < signatures->size(); ++i) {
auto* signature_def = signatures->Get(i);
out_stream << "Signature#" << i << " key: '"
<< signature_def->signature_key()->str() << "'\n";
out_stream << "- Subgraph: "
<< subgraph_str(signature_def->subgraph_index()) << "\n";
out_stream << "- Inputs: \n";
for (int j = 0; j < signature_def->inputs()->size(); ++j) {
auto* input = signature_def->inputs()->Get(j);
out_stream << " '" << input->name()->str() << "' : "
<< tensor_str(input->tensor_index(),
signature_def->subgraph_index())
<< "\n";
}
out_stream << "- Outputs: \n";
for (int j = 0; j < signature_def->outputs()->size(); ++j) {
auto* output = signature_def->outputs()->Get(j);
out_stream << " '" << output->name()->str() << "' : "
<< tensor_str(output->tensor_index(),
signature_def->subgraph_index())
<< "\n";
}
out_stream << "\n";
}
}
// Dump the statistics of the given TFLite flatbuffer model. It's printed at the
// end of the analyzer output.
void dump_model_stats(std::stringstream& out_stream,
const ::tflite::Model* model, size_t model_size,
ModelStats* stats) {
size_t total_buffer_size = 0;
size_t total_zero_buffer_size = 0;
auto* buffers = model->buffers();
for (int i = 0; i < buffers->size(); ++i) {
const tflite::Buffer* buffer = buffers->Get(i);
if (buffer->data() == nullptr) {
continue;
}
bool is_all_zeros = true;
const unsigned char* data = buffer->data()->data();
for (int j = 0; j < buffer->data()->size(); ++j) {
if (data[j] != 0) {
is_all_zeros = false;
break;
}
}
if (is_all_zeros) {
total_zero_buffer_size += buffer->data()->size();
}
total_buffer_size += buffer->data()->size();
}
out_stream << kSectionSplitter;
char temp[2048];
snprintf(temp, sizeof(temp), "%24s: %10zu bytes\n", "Model size", model_size);
out_stream << temp;
snprintf(
temp, sizeof(temp), "%24s: %10zu bytes (%05.2f %%)\n",
"Non-data buffer size", model_size - total_buffer_size,
(static_cast<float>(model_size - total_buffer_size) / model_size * 100));
out_stream << temp;
snprintf(temp, sizeof(temp), "%24s: %10zu bytes (%05.2f %%)\n",
"Total data buffer size", total_buffer_size,
(static_cast<float>(total_buffer_size) / model_size * 100));
out_stream << temp;
if (model->subgraphs()->size() > 1) {
for (int i = 0; i < model->subgraphs()->size(); ++i) {
float subgraph_buffer_ratio =
static_cast<float>(stats->buffer_usage[i]) / model_size * 100;
snprintf(temp, sizeof(temp),
" - %-12s: %10zu bytes (%05.2f %%)\n",
subgraph_str(i).c_str(), stats->buffer_usage[i],
subgraph_buffer_ratio);
out_stream << temp;
}
}
float zero_buffer_ratio =
static_cast<float>(total_zero_buffer_size) / model_size * 100;
snprintf(temp, sizeof(temp), "%24s: %10zu bytes (%05.2f %%)\n",
"(Zero value buffers)", total_zero_buffer_size, zero_buffer_ratio);
out_stream << temp;
out_stream
<< "\n"
<< "* Buffers of TFLite model are mostly used for constant tensors.\n";
out_stream << " And zero value buffers are buffers filled with zeros.\n";
if (zero_buffer_ratio > kThreshold_zero_buffer_ratio) {
out_stream << " (Consider use "
"`converter._experimental_unfold_large_splat_constant` "
"to save the model size.)\n";
}
out_stream << " Non-data buffers area are used to store operators, "
"subgraphs and etc.\n";
out_stream << " You can find more details from "
"https://github.com/tensorflow/tensorflow/blob/master/"
"tensorflow/lite/schema/schema.fbs\n";
}
} // namespace
class StreamErrorReporter : public ErrorReporter {
public:
explicit StreamErrorReporter(std::stringstream* out_stream)
: out_stream_(out_stream) {}
int Report(const char* format, va_list args) override {
char buffer[1024];
int size = vsnprintf(buffer, sizeof(buffer), format, args);
*out_stream_ << buffer;
return size;
}
private:
std::stringstream* out_stream_;
};
std::string get_printable_string(const std::string& src) {
std::stringstream out;
for (auto it = src.begin(); it != src.end(); ++it) {
out << ((*it == '\n' || isprint(*it)) ? *it : '.');
}
return out.str();
}
std::string model_analyzer(const std::string& model_file_or_buffer,
bool input_is_filepath,
bool check_gpu_compatibility) {
std::stringstream out_stream;
StreamErrorReporter error_reporter(&out_stream);
std::unique_ptr<FlatBufferModel> fb_model;
if (input_is_filepath) {
fb_model = FlatBufferModel::BuildFromFile(model_file_or_buffer.c_str(),
&error_reporter);
if (!fb_model) {
out_stream << "Failed to mmap model " << model_file_or_buffer;
return out_stream.str();
}
} else {
fb_model = FlatBufferModel::BuildFromBuffer(model_file_or_buffer.c_str(),
model_file_or_buffer.size(),
&error_reporter);
if (!fb_model) {
out_stream << "Failed to mmap the given model buffer.";
return out_stream.str();
}
}
const ::tflite::Model* model = fb_model->GetModel();
auto* subgraphs = model->subgraphs();
ModelStats stats;
stats.buffer_usage.resize(subgraphs->size());
dump_model_summary(out_stream, model);
bool model_is_gpu_compatible = true;
for (int i = 0; i < subgraphs->size(); ++i) {
std::vector<int> gpu_incompatible_nodes;
const SubGraph* subgraph = subgraphs->Get(i);
out_stream << subgraph_str(i);
if (subgraph->name()) {
out_stream << " " << subgraph->name()->str();
}
out_stream << "(";
dump_tensor_list(out_stream, subgraph->inputs(), i);
out_stream << ") -> [";
dump_tensor_list(out_stream, subgraph->outputs(), i);
out_stream << "]\n";
if (subgraph->operators()) {
for (int j = 0; j < subgraph->operators()->size(); ++j) {
const Operator* op = subgraph->operators()->Get(j);
const OperatorCode* op_code =
model->operator_codes()->Get(op->opcode_index());
out_stream << " "; // indents for operators
dump_node(out_stream, /*node_no=*/j, op_code, op, i, model);
}
}
if (!gpu_incompatible_nodes.empty()) {
model_is_gpu_compatible = false;
out_stream << "\nGPU COMPATIBILITY WARNING: Subgraph#" << i
<< " has GPU delegate compatibility issues at nodes "
<< absl::StrJoin(gpu_incompatible_nodes, ", ")
<< " on TFLite runtime version " << TFLITE_VERSION_STRING
<< "\n";
}
// Dump Subgraph Tensors.
out_stream << "\nTensors of " << subgraph_str(i) << "\n";
auto tensors = subgraph->tensors();
if (tensors) {
for (int j = 0; j < tensors->size(); ++j) {
auto tensor = tensors->Get(j);
out_stream << " "; // indents for tensors
dump_tensor_detail(out_stream, tensor, j, i, model, &stats);
}
}
out_stream << "\n";
}
if (check_gpu_compatibility && model_is_gpu_compatible) {
out_stream
<< "\nYour model looks compatible with GPU delegate"
<< " on TFLite runtime version " << TFLITE_VERSION_STRING
<< ".\nThis does not guarantee that your model will work well with GPU"
" delegate because there could still be runtime "
"incompatibililties.\n";
}
dump_model_signature_defs(out_stream, model);
dump_model_stats(out_stream, model, fb_model->allocation()->bytes(), &stats);
return get_printable_string(out_stream.str());
}
} // namespace tflite
tensorflow/lite/python/analyzer_wrapper/model_analyzer.h
+30
View File
@@ -0,0 +1,30 @@
/* Copyright 2021 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_PYTHON_ANALYZER_WRAPPER_MODEL_ANALYZER_H_
#define TENSORFLOW_LITE_PYTHON_ANALYZER_WRAPPER_MODEL_ANALYZER_H_
#include <string>
namespace tflite {
// Returns a brief dump of the given TFLite file or model.
// It examines the model file itself without instantiating TFLite interpreters.
std::string model_analyzer(const std::string& model_file_or_buffer,
bool input_is_filepath,
bool check_gpu_compatibility);
} // namespace tflite
#endif // TENSORFLOW_LITE_PYTHON_ANALYZER_WRAPPER_MODEL_ANALYZER_H_
tensorflow/lite/python/authoring/BUILD
+34
View File
@@ -0,0 +1,34 @@
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
load("//tensorflow:tensorflow.bzl", "py_test")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
licenses = ["notice"],
)
py_library(
name = "authoring",
srcs = [
"authoring.py",
],
strict_deps = True,
visibility = ["//visibility:public"],
deps = [
"//tensorflow/compiler/mlir/lite/metrics:converter_error_data_proto_py",
"//tensorflow/lite/python:convert",
"//tensorflow/lite/python:lite",
"//tensorflow/python/util:tf_export",
],
)
py_test(
name = "authoring_test",
srcs = ["authoring_test.py"],
strict_deps = True,
deps = [
":authoring",
#internal proto upb dep
"//tensorflow:tensorflow_py",
"//tensorflow/lite/python:lite",
],
)
tensorflow/lite/python/authoring/authoring.py
+301
View File
@@ -0,0 +1,301 @@
# Copyright 2021 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 Authoring tool package for TFLite compatibility.
WARNING: The package is experimental and subject to change.
This package provides a way to check TFLite compatibility at model authoring
time.
Example:
@tf.lite.experimental.authoring.compatible
@tf.function(input_signature=[
tf.TensorSpec(shape=[None], dtype=tf.float32)
])
def f(x):
return tf.cosh(x)
result = f(tf.constant([0.0]))
> COMPATIBILITY WARNING: op 'tf.Cosh' require(s) "Select TF Ops" for model
> conversion for TensorFlow Lite.
> Op: tf.Cosh
> - tensorflow/python/framework/op_def_library.py:xxx
> - tensorflow/python/ops/gen_math_ops.py:xxx
> - simple_authoring.py:xxx
"""
import functools
from tensorflow.compiler.mlir.lite.metrics import converter_error_data_pb2
# pylint: disable=g-import-not-at-top
from tensorflow.lite.python import convert
from tensorflow.lite.python import lite
from tensorflow.python.util.tf_export import tf_export as _tf_export
_CUSTOM_OPS_HDR = "Custom ops: "
_TF_OPS_HDR = "TF Select ops: "
_AUTHORING_ERROR_HDR = "COMPATIBILITY ERROR"
_AUTHORING_WARNING_HDR = "COMPATIBILITY WARNING"
_FUNC_GRAPH_SRC_PATH = "tensorflow/python/framework/func_graph.py"
class CompatibilityError(Exception):
"""Raised when an error occurs with TFLite compatibility."""
pass
class _Compatible:
"""A decorator class to check TFLite compatibility created by `lite.experimental.authoring.compatible`."""
def __init__(self,
target,
converter_target_spec=None,
converter_allow_custom_ops=None,
raise_exception=False):
"""Initialize the decorator object.
Here is the description of the object variables.
- _func : decorated function.
- _obj_func : for class object, we need to use this object to provide `self`
instance as 1 first argument.
- _verified : whether the compatibility is checked or not.
Args:
target: decorated function.
converter_target_spec : target_spec of TFLite converter parameter.
converter_allow_custom_ops : allow_custom_ops of TFLite converter
parameter.
raise_exception : to raise an exception on compatibility issues.
User need to use get_compatibility_log() to check details.
"""
functools.update_wrapper(self, target)
self._func = target
self._obj_func = None
self._verified = False
self._log_messages = []
self._raise_exception = raise_exception
self._converter_target_spec = converter_target_spec
self._converter_allow_custom_ops = converter_allow_custom_ops
def __get__(self, instance, cls):
"""A Python descriptor interface."""
self._obj_func = self._func.__get__(instance, cls)
return self
def _get_func(self):
"""Returns decorated function object.
For a class method, use self._obj_func to provide `self` instance.
"""
if self._obj_func is not None:
return self._obj_func
else:
return self._func
def __call__(self, *args, **kwargs): # pylint: disable=g-doc-args
"""Calls decorated function object.
Also verifies if the function is compatible with TFLite.
Returns:
A execution result of the decorated function.
"""
if not self._verified:
model = self._get_func()
concrete_func = model.get_concrete_function(*args, **kwargs)
converter = lite.TFLiteConverterV2.from_concrete_functions(
[concrete_func], model)
# Set provided converter parameters
if self._converter_target_spec is not None:
converter.target_spec = self._converter_target_spec
if self._converter_allow_custom_ops is not None:
converter.allow_custom_ops = self._converter_allow_custom_ops
try:
converter.convert()
except convert.ConverterError as err:
self._decode_error(err)
finally:
self._verified = True
return self._get_func()(*args, **kwargs)
def get_concrete_function(self, *args, **kwargs):
"""Returns a concrete function of the decorated function."""
return self._get_func().get_concrete_function(*args, **kwargs)
def _get_location_string(self, location):
"""Dump location of ConveterError.errors.location."""
callstack = []
for single_call in reversed(location.call):
if (location.type ==
converter_error_data_pb2.ConverterErrorData.CALLSITELOC):
callstack.append(
f" - {single_call.source.filename}:{single_call.source.line}")
else:
callstack.append(str(single_call))
callstack_dump = "\n".join(callstack)
return callstack_dump
def _dump_error_details(self, ops, locations):
"""Dump the list of ops and locations."""
for i in range(0, len(ops)):
callstack_dump = self._get_location_string(locations[i])
err_string = f"Op: {ops[i]}\n{callstack_dump}\n"
self._log(err_string)
def _decode_error_legacy(self, err):
"""Parses the given legacy ConverterError for OSS."""
for line in str(err).splitlines():
# Check custom op usage error.
if line.startswith(_CUSTOM_OPS_HDR):
custom_ops = line[len(_CUSTOM_OPS_HDR):]
err_string = (
f"{_AUTHORING_ERROR_HDR}: op '{custom_ops}' is(are) not natively "
"supported by TensorFlow Lite. You need to provide a custom "
"operator. https://www.tensorflow.org/lite/guide/ops_custom")
self._log(err_string)
# Check TensorFlow op usage error.
elif line.startswith(_TF_OPS_HDR):
tf_ops = line[len(_TF_OPS_HDR):]
err_string = (
f"{_AUTHORING_WARNING_HDR}: op '{tf_ops}' require(s) \"Select TF "
"Ops\" for model conversion for TensorFlow Lite. "
"https://www.tensorflow.org/lite/guide/ops_select")
self._log(err_string)
def _decode_converter_error(self, err):
"""Parses the given ConverterError which has detailed error information."""
custom_ops = []
custom_ops_location = []
tf_ops = []
tf_ops_location = []
gpu_not_compatible_ops = []
for err in err.errors:
# Check custom op usage error.
if err.error_code == converter_error_data_pb2.ConverterErrorData.ERROR_NEEDS_CUSTOM_OPS:
custom_ops.append(err.operator.name)
custom_ops_location.append(err.location)
# Check TensorFlow op usage error.
elif err.error_code == converter_error_data_pb2.ConverterErrorData.ERROR_NEEDS_FLEX_OPS:
tf_ops.append(err.operator.name)
tf_ops_location.append(err.location)
# Check GPU delegate compatibility error.
elif err.error_code == converter_error_data_pb2.ConverterErrorData.ERROR_GPU_NOT_COMPATIBLE:
gpu_not_compatible_ops.append(err.operator.name)
# Log the first line of ConveterError.errors.error_message only
# since the seond line is "Error code: xxxx"
self._log(err.error_message.splitlines()[0])
self._log(self._get_location_string(err.location) + "\n")
else:
# Log other errors.
self._log(f"{_AUTHORING_ERROR_HDR}: {err.error_message}")
self._log(self._get_location_string(err.location) + "\n")
if custom_ops:
custom_ops_str = ", ".join(sorted(custom_ops))
err_string = (
f"{_AUTHORING_ERROR_HDR}: op '{custom_ops_str}' is(are) not natively "
"supported by TensorFlow Lite. You need to provide a custom "
"operator. https://www.tensorflow.org/lite/guide/ops_custom")
self._log(err_string)
self._dump_error_details(custom_ops, custom_ops_location)
if tf_ops:
tf_ops_str = ", ".join(sorted(tf_ops))
err_string = (
f"{_AUTHORING_WARNING_HDR}: op '{tf_ops_str}' require(s) \"Select TF"
" Ops\" for model conversion for TensorFlow Lite. "
"https://www.tensorflow.org/lite/guide/ops_select")
self._log(err_string)
self._dump_error_details(tf_ops, tf_ops_location)
if gpu_not_compatible_ops:
not_compatible_ops_str = ", ".join(sorted(gpu_not_compatible_ops))
err_string = (
f"{_AUTHORING_WARNING_HDR}: op '{not_compatible_ops_str}' aren't "
"compatible with TensorFlow Lite GPU delegate. "
"https://www.tensorflow.org/lite/performance/gpu")
self._log(err_string)
def _decode_error(self, err):
"""Parses the given ConverterError and generates compatibility warnings."""
if hasattr(err, "errors"):
self._decode_converter_error(err)
else:
self._decode_error_legacy(err)
if self._raise_exception and self._log_messages:
raise CompatibilityError(f"CompatibilityException at {repr(self._func)}")
def _log(self, message):
"""Log and print authoring warning / error message."""
self._log_messages.append(message)
print(message)
def get_compatibility_log(self):
"""Returns list of compatibility log messages.
WARNING: This method should only be used for unit tests.
Returns:
The list of log messages by the recent compatibility check.
Raises:
RuntimeError: when the compatibility was NOT checked.
"""
if not self._verified:
raise RuntimeError("target compatibility isn't verified yet")
return self._log_messages
@_tf_export("lite.experimental.authoring.compatible")
def compatible(target=None, converter_target_spec=None, **kwargs):
"""Wraps `tf.function` into a callable function with TFLite compatibility checking.
Example:
```python
@tf.lite.experimental.authoring.compatible
@tf.function(input_signature=[
tf.TensorSpec(shape=[None], dtype=tf.float32)
])
def f(x):
return tf.cosh(x)
result = f(tf.constant([0.0]))
# COMPATIBILITY WARNING: op 'tf.Cosh' require(s) "Select TF Ops" for model
# conversion for TensorFlow Lite.
# Op: tf.Cosh
# - tensorflow/python/framework/op_def_library.py:748
# - tensorflow/python/ops/gen_math_ops.py:2458
# - <stdin>:6
```
WARNING: Experimental interface, subject to change.
Args:
target: A `tf.function` to decorate.
converter_target_spec : target_spec of TFLite converter parameter.
**kwargs: The keyword arguments of the decorator class _Compatible.
Returns:
A callable object of `tf.lite.experimental.authoring._Compatible`.
"""
if target is None:
def wrapper(target):
return _Compatible(target, converter_target_spec, **kwargs)
return wrapper
else:
return _Compatible(target, converter_target_spec, **kwargs)
tensorflow/lite/python/authoring/authoring_test.py
+298
View File
@@ -0,0 +1,298 @@
# Copyright 2021 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.
# ==============================================================================
"""Unit tests for authoring package."""
# pylint: disable=g-direct-tensorflow-import
import tensorflow as tf
from tensorflow.lite.python.authoring import authoring
class TFLiteAuthoringTest(tf.test.TestCase):
def test_simple_cosh(self):
@authoring.compatible
@tf.function(input_signature=[
tf.TensorSpec(shape=[None], dtype=tf.float32)
])
def f(x):
return tf.cosh(x)
result = f(tf.constant([0.0]))
log_messages = f.get_compatibility_log()
self.assertEqual(result, tf.constant([1.0]))
self.assertIn(
"COMPATIBILITY WARNING: op 'tf.Cosh' require(s) \"Select TF Ops\" for "
"model conversion for TensorFlow Lite. "
"https://www.tensorflow.org/lite/guide/ops_select", log_messages)
# Check the op location ends with filename of the this test.
self.assertIn("authoring_test.py", log_messages[-1])
def test_simple_cosh_raises_CompatibilityError(self):
@authoring.compatible(raise_exception=True)
@tf.function(input_signature=[
tf.TensorSpec(shape=[None], dtype=tf.float32)
])
def f(x):
return tf.cosh(x)
# Check if the CompatibilityError exception is raised.
with self.assertRaises(authoring.CompatibilityError):
result = f(tf.constant([0.0]))
del result
log_messages = f.get_compatibility_log()
self.assertIn(
"COMPATIBILITY WARNING: op 'tf.Cosh' require(s) \"Select TF Ops\" for "
"model conversion for TensorFlow Lite. "
"https://www.tensorflow.org/lite/guide/ops_select", log_messages)
def test_flex_compatibility(self):
@authoring.compatible
@tf.function(input_signature=[
tf.TensorSpec(shape=[3, 3, 3, 3, 3], dtype=tf.float32)
])
def f(inp):
tanh = tf.math.tanh(inp)
conv3d = tf.nn.conv3d(
tanh,
tf.ones([3, 3, 3, 3, 3]),
strides=[1, 1, 1, 1, 1],
padding="SAME")
erf = tf.math.erf(conv3d)
output = tf.math.tanh(erf)
return output
f(tf.ones(shape=(3, 3, 3, 3, 3), dtype=tf.float32))
log_messages = f.get_compatibility_log()
self.assertIn(
"COMPATIBILITY WARNING: op 'tf.Erf' require(s) \"Select TF Ops\" for "
"model conversion for TensorFlow Lite. "
"https://www.tensorflow.org/lite/guide/ops_select", log_messages)
def test_compatibility_error_generic(self):
@authoring.compatible
@tf.function
def f():
dataset = tf.data.Dataset.range(3)
dataset = dataset.shuffle(3, reshuffle_each_iteration=True)
return dataset
f()
log_messages = f.get_compatibility_log()
self.assertIn(
"COMPATIBILITY ERROR: op 'tf.DummySeedGenerator, tf.RangeDataset,"
" tf.ShuffleDatasetV3' is(are) not natively supported by TensorFlow"
" Lite. You need to provide a custom operator."
" https://www.tensorflow.org/lite/guide/ops_custom",
log_messages,
)
def test_compatibility_error_custom(self):
target_spec = tf.lite.TargetSpec()
target_spec.supported_ops = [
tf.lite.OpsSet.TFLITE_BUILTINS,
tf.lite.OpsSet.SELECT_TF_OPS,
]
@authoring.compatible(converter_target_spec=target_spec)
@tf.function
def f():
dataset = tf.data.Dataset.range(3)
dataset = dataset.shuffle(3, reshuffle_each_iteration=True)
return dataset
f()
log_messages = f.get_compatibility_log()
self.assertIn(
"COMPATIBILITY ERROR: op 'tf.DummySeedGenerator, tf.RangeDataset, "
"tf.ShuffleDatasetV3' is(are) not natively supported by "
"TensorFlow Lite. You need to provide a custom operator. "
"https://www.tensorflow.org/lite/guide/ops_custom", log_messages)
def test_simple_variable(self):
external_var = tf.Variable(1.0)
@authoring.compatible
@tf.function(input_signature=[
tf.TensorSpec(shape=[None], dtype=tf.float32)
])
def f(x):
return x * external_var
result = f(tf.constant(2.0, shape=(1)))
log_messages = f.get_compatibility_log()
self.assertEqual(result, tf.constant([2.0]))
self.assertEmpty(log_messages)
def test_class_method(self):
class Model(tf.Module):
@authoring.compatible
@tf.function(input_signature=[
tf.TensorSpec(shape=[None], dtype=tf.float32)
])
def eval(self, x):
return tf.cosh(x)
m = Model()
result = m.eval(tf.constant([0.0]))
log_messages = m.eval.get_compatibility_log()
self.assertEqual(result, tf.constant([1.0]))
self.assertIn(
"COMPATIBILITY WARNING: op 'tf.Cosh' require(s) \"Select TF Ops\" for "
"model conversion for TensorFlow Lite. "
"https://www.tensorflow.org/lite/guide/ops_select", log_messages)
def test_decorated_function_type(self):
@authoring.compatible
@tf.function(input_signature=[
tf.TensorSpec(shape=[None], dtype=tf.float32)
])
def func(x):
return tf.cos(x)
result = func(tf.constant([0.0]))
self.assertEqual(result, tf.constant([1.0]))
# Check if the decorator keeps __name__ attribute.
self.assertEqual(func.__name__, "func")
# Check if the decorator works with get_concrete_function method.
converter = tf.lite.TFLiteConverter.from_concrete_functions(
[func.get_concrete_function()], func)
converter.convert()
def test_decorated_class_method_type(self):
class Model(tf.Module):
@authoring.compatible
@tf.function(input_signature=[
tf.TensorSpec(shape=[None], dtype=tf.float32)
])
def eval(self, x):
return tf.cos(x)
m = Model()
result = m.eval(tf.constant([0.0]))
self.assertEqual(result, tf.constant([1.0]))
# Check if the decorator keeps __name__ attribute.
self.assertEqual(m.eval.__name__, "eval")
# Check if the decorator works with get_concrete_function method.
converter = tf.lite.TFLiteConverter.from_concrete_functions(
[m.eval.get_concrete_function()], m)
converter.convert()
def test_simple_cosh_multiple(self):
@authoring.compatible
@tf.function(input_signature=[
tf.TensorSpec(shape=[None], dtype=tf.float32)
])
def f(x):
return tf.cosh(x)
f(tf.constant([1.0]))
f(tf.constant([2.0]))
f(tf.constant([3.0]))
warning_messages = f.get_compatibility_log()
# Test if compatiblility checks happens only once.
# The number of warning_messages will be 2 by op location detail.
self.assertEqual(2, len(warning_messages))
def test_user_tf_ops_all_filtered(self):
target_spec = tf.lite.TargetSpec()
target_spec.supported_ops = [
tf.lite.OpsSet.TFLITE_BUILTINS,
tf.lite.OpsSet.SELECT_TF_OPS,
]
target_spec.experimental_select_user_tf_ops = [
"RangeDataset", "DummySeedGenerator", "ShuffleDatasetV3"
]
@authoring.compatible(converter_target_spec=target_spec)
@tf.function
def f():
dataset = tf.data.Dataset.range(3)
dataset = dataset.shuffle(3, reshuffle_each_iteration=True)
return dataset
f()
log_messages = f.get_compatibility_log()
self.assertEmpty(log_messages)
def test_user_tf_ops_partial_filtered(self):
target_spec = tf.lite.TargetSpec()
target_spec.supported_ops = [
tf.lite.OpsSet.TFLITE_BUILTINS,
tf.lite.OpsSet.SELECT_TF_OPS,
]
target_spec.experimental_select_user_tf_ops = [
"DummySeedGenerator"
]
@authoring.compatible(converter_target_spec=target_spec)
@tf.function
def f():
dataset = tf.data.Dataset.range(3)
dataset = dataset.shuffle(3, reshuffle_each_iteration=True)
return dataset
f()
log_messages = f.get_compatibility_log()
self.assertIn(
"COMPATIBILITY ERROR: op 'tf.RangeDataset, tf.ShuffleDatasetV3' is(are) "
"not natively supported by TensorFlow Lite. You need to provide a "
"custom operator. https://www.tensorflow.org/lite/guide/ops_custom",
log_messages)
def test_allow_custom_ops(self):
target_spec = tf.lite.TargetSpec()
target_spec.supported_ops = [
tf.lite.OpsSet.TFLITE_BUILTINS,
tf.lite.OpsSet.SELECT_TF_OPS,
]
@authoring.compatible(
converter_allow_custom_ops=True, converter_target_spec=target_spec)
@tf.function
def f():
dataset = tf.data.Dataset.range(3)
dataset = dataset.shuffle(3, reshuffle_each_iteration=True)
return dataset
f()
log_messages = f.get_compatibility_log()
self.assertEmpty(log_messages)
def test_gpu_compatible(self):
target_spec = tf.lite.TargetSpec()
target_spec.supported_ops = [
tf.lite.OpsSet.TFLITE_BUILTINS,
]
target_spec.experimental_supported_backends = ["GPU"]
@authoring.compatible(converter_target_spec=target_spec)
@tf.function(
input_signature=[tf.TensorSpec(shape=[4, 4], dtype=tf.float32)])
def func(x):
return tf.cos(x)
func(tf.ones(shape=(4, 4), dtype=tf.float32))
log_messages = func.get_compatibility_log()
self.assertEmpty(log_messages)
if __name__ == "__main__":
tf.test.main()
tensorflow/lite/python/convert.py
+1101
View File
File diff suppressed because it is too large Load Diff
tensorflow/lite/python/convert_file_to_c_source.py
+68
View File
@@ -0,0 +1,68 @@
# Copyright 2018 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 TFLite model to a TFLite Micro model (C++ Source)."""
from absl import app
from absl import flags
from tensorflow.lite.python import util
FLAGS = flags.FLAGS
flags.DEFINE_string("input_tflite_file", None,
"Full path name to the input TFLite model file.")
flags.DEFINE_string(
"output_source_file", None,
"Full path name to the output TFLite Micro model (C++ Source) file).")
flags.DEFINE_string("output_header_file", None,
"Full filepath of the output C header file.")
flags.DEFINE_string("array_variable_name", None,
"Name to use for the C data array variable.")
flags.DEFINE_integer("line_width", 80, "Width to use for formatting.")
flags.DEFINE_string("include_guard", None,
"Name to use for the C header include guard.")
flags.DEFINE_string("include_path", None,
"Optional path to include in generated source file.")
flags.DEFINE_boolean(
"use_tensorflow_license", False,
"Whether to prefix the generated files with the TF Apache2 license.")
flags.mark_flag_as_required("input_tflite_file")
flags.mark_flag_as_required("output_source_file")
flags.mark_flag_as_required("output_header_file")
flags.mark_flag_as_required("array_variable_name")
def main(_):
with open(FLAGS.input_tflite_file, "rb") as input_handle:
input_data = input_handle.read()
source, header = util.convert_bytes_to_c_source(
data=input_data,
array_name=FLAGS.array_variable_name,
max_line_width=FLAGS.line_width,
include_guard=FLAGS.include_guard,
include_path=FLAGS.include_path,
use_tensorflow_license=FLAGS.use_tensorflow_license)
with open(FLAGS.output_source_file, "w") as source_handle:
source_handle.write(source)
with open(FLAGS.output_header_file, "w") as header_handle:
header_handle.write(header)
if __name__ == "__main__":
app.run(main)
tensorflow/lite/python/convert_file_to_c_source_test.sh
Executable
+87
View File
@@ -0,0 +1,87 @@
#!/bin/bash
# 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.
# ==============================================================================
#
# Bash unit tests for the TensorFlow Lite Micro project generator.
set -e
INPUT_FILE=${TEST_TMPDIR}/input.tflite
printf "\x00\x01\x02\x03" > ${INPUT_FILE}
OUTPUT_SOURCE_FILE=${TEST_TMPDIR}/output_source.cc
OUTPUT_HEADER_FILE=${TEST_TMPDIR}/output_header.h
# Needed for copybara compatibility.
SCRIPT_BASE_DIR=/org_"tensor"flow
${TEST_SRCDIR}${SCRIPT_BASE_DIR}/tensorflow/lite/python/convert_file_to_c_source \
--input_tflite_file="${INPUT_FILE}" \
--output_source_file="${OUTPUT_SOURCE_FILE}" \
--output_header_file="${OUTPUT_HEADER_FILE}" \
--array_variable_name="g_some_array" \
--line_width=80 \
--include_guard="SOME_GUARD_H_" \
--include_path="some/guard.h" \
--use_tensorflow_license=True
if ! grep -q 'const unsigned char g_some_array' ${OUTPUT_SOURCE_FILE}; then
echo "ERROR: No array found in output '${OUTPUT_SOURCE_FILE}'"
exit 1
fi
if ! grep -q '0x00, 0x01, 0x02, 0x03' ${OUTPUT_SOURCE_FILE}; then
echo "ERROR: No array values found in output '${OUTPUT_SOURCE_FILE}'"
exit 1
fi
if ! grep -q 'const int g_some_array_len = 4;' ${OUTPUT_SOURCE_FILE}; then
echo "ERROR: No array length found in output '${OUTPUT_SOURCE_FILE}'"
exit 1
fi
if ! grep -q 'The TensorFlow Authors. All Rights Reserved' ${OUTPUT_SOURCE_FILE}; then
echo "ERROR: No license found in output '${OUTPUT_SOURCE_FILE}'"
exit 1
fi
if ! grep -q '\#include "some/guard\.h"' ${OUTPUT_SOURCE_FILE}; then
echo "ERROR: No include found in output '${OUTPUT_SOURCE_FILE}'"
exit 1
fi
if ! grep -q '#ifndef SOME_GUARD_H_' ${OUTPUT_HEADER_FILE}; then
echo "ERROR: No include guard found in output '${OUTPUT_HEADER_FILE}'"
exit 1
fi
if ! grep -q 'extern const unsigned char g_some_array' ${OUTPUT_HEADER_FILE}; then
echo "ERROR: No array found in output '${OUTPUT_HEADER_FILE}'"
exit 1
fi
if ! grep -q 'extern const int g_some_array_len;' ${OUTPUT_HEADER_FILE}; then
echo "ERROR: No array length found in output '${OUTPUT_HEADER_FILE}'"
exit 1
fi
if ! grep -q 'The TensorFlow Authors. All Rights Reserved' ${OUTPUT_HEADER_FILE}; then
echo "ERROR: No license found in output '${OUTPUT_HEADER_FILE}'"
exit 1
fi
echo
echo "SUCCESS: convert_file_to_c_source test PASSED"
tensorflow/lite/python/convert_phase.py
+219
View File
@@ -0,0 +1,219 @@
# Copyright 2021 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.
# ==============================================================================
"""Utilities for collecting TFLite metrics."""
import collections
import enum
import functools
from typing import Text
from tensorflow.compiler.mlir.lite.metrics import converter_error_data_pb2
from tensorflow.lite.python.metrics import metrics
class Component(enum.Enum):
"""Enum class defining name of the converter components."""
# Validate the given input and prepare and optimize TensorFlow Model.
PREPARE_TF_MODEL = "PREPARE_TF_MODEL"
# Convert to TFLite model format.
CONVERT_TF_TO_TFLITE_MODEL = "CONVERT_TF_TO_TFLITE_MODEL"
# RUN quantization and sparsification.
OPTIMIZE_TFLITE_MODEL = "OPTIMIZE_TFLITE_MODEL"
SubComponentItem = collections.namedtuple("SubComponentItem",
["name", "component"])
class SubComponent(SubComponentItem, enum.Enum):
"""Enum class defining name of the converter subcomponents.
This enum only defines the subcomponents in Python, there might be more
subcomponents defined in C++.
"""
def __str__(self):
return self.value.name
@property
def name(self):
return self.value.name
@property
def component(self):
return self.value.component
# The subcomponent name is unspecified.
UNSPECIFIED = SubComponentItem("UNSPECIFIED", None)
# Valid the given input and parameters.
VALIDATE_INPUTS = SubComponentItem("VALIDATE_INPUTS",
Component.PREPARE_TF_MODEL)
# Load GraphDef from SavedModel.
LOAD_SAVED_MODEL = SubComponentItem("LOAD_SAVED_MODEL",
Component.PREPARE_TF_MODEL)
# Convert a SavedModel to frozen graph.
FREEZE_SAVED_MODEL = SubComponentItem("FREEZE_SAVED_MODEL",
Component.PREPARE_TF_MODEL)
# Save a Keras model to SavedModel.
CONVERT_KERAS_TO_SAVED_MODEL = SubComponentItem(
"CONVERT_KERAS_TO_SAVED_MODEL", Component.PREPARE_TF_MODEL)
# Save Concrete functions to SavedModel.
CONVERT_CONCRETE_FUNCTIONS_TO_SAVED_MODEL = SubComponentItem(
"CONVERT_CONCRETE_FUNCTIONS_TO_SAVED_MODEL", Component.PREPARE_TF_MODEL)
# Convert a Keras model to a frozen graph.
FREEZE_KERAS_MODEL = SubComponentItem("FREEZE_KERAS_MODEL",
Component.PREPARE_TF_MODEL)
# Replace all the variables with constants in a ConcreteFunction.
FREEZE_CONCRETE_FUNCTION = SubComponentItem("FREEZE_CONCRETE_FUNCTION",
Component.PREPARE_TF_MODEL)
# Run grappler optimization.
OPTIMIZE_TF_MODEL = SubComponentItem("OPTIMIZE_TF_MODEL",
Component.PREPARE_TF_MODEL)
# Convert using the old TOCO converter.
CONVERT_GRAPHDEF_USING_DEPRECATED_CONVERTER = SubComponentItem(
"CONVERT_GRAPHDEF_USING_DEPRECATED_CONVERTER",
Component.CONVERT_TF_TO_TFLITE_MODEL)
# Convert a GraphDef to TFLite model.
CONVERT_GRAPHDEF = SubComponentItem("CONVERT_GRAPHDEF",
Component.CONVERT_TF_TO_TFLITE_MODEL)
# Convert a SavedModel to TFLite model.
CONVERT_SAVED_MODEL = SubComponentItem("CONVERT_SAVED_MODEL",
Component.CONVERT_TF_TO_TFLITE_MODEL)
# Convert a Jax HLO to TFLite model.
CONVERT_JAX_HLO = SubComponentItem("CONVERT_JAX_HLO",
Component.CONVERT_TF_TO_TFLITE_MODEL)
# Do quantization by the deprecated quantizer.
QUANTIZE_USING_DEPRECATED_QUANTIZER = SubComponentItem(
"QUANTIZE_USING_DEPRECATED_QUANTIZER", Component.OPTIMIZE_TFLITE_MODEL)
# Do calibration.
CALIBRATE = SubComponentItem("CALIBRATE", Component.OPTIMIZE_TFLITE_MODEL)
# Do quantization by MLIR.
QUANTIZE = SubComponentItem("QUANTIZE", Component.OPTIMIZE_TFLITE_MODEL)
# Do sparsification by MLIR.
SPARSIFY = SubComponentItem("SPARSIFY", Component.OPTIMIZE_TFLITE_MODEL)
class ConverterError(Exception):
"""Raised when an error occurs during model conversion."""
def __init__(self, message):
super(ConverterError, self).__init__(message)
self.errors = []
self._parse_error_message(message)
def append_error(self,
error_data: converter_error_data_pb2.ConverterErrorData):
self.errors.append(error_data)
def _parse_error_message(self, message):
"""If the message matches a pattern, assigns the associated error code.
It is difficult to assign an error code to some errrors in MLIR side, Ex:
errors thrown by other components than TFLite or not using mlir::emitError.
This function try to detect them by the error message and assign the
corresponding error code.
Args:
message: The error message of this exception.
"""
error_code_mapping = {
"Failed to functionalize Control Flow V1 ops. Consider using Control "
"Flow V2 ops instead. See https://www.tensorflow.org/api_docs/python/"
"tf/compat/v1/enable_control_flow_v2.":
converter_error_data_pb2.ConverterErrorData
.ERROR_UNSUPPORTED_CONTROL_FLOW_V1,
}
for pattern, error_code in error_code_mapping.items():
if pattern in message:
error_data = converter_error_data_pb2.ConverterErrorData()
error_data.error_message = message
error_data.error_code = error_code
self.append_error(error_data)
return
def convert_phase(component, subcomponent=SubComponent.UNSPECIFIED):
"""The decorator to identify converter component and subcomponent.
Args:
component: Converter component name.
subcomponent: Converter subcomponent name.
Returns:
Forward the result from the wrapped function.
Raises:
ValueError: if component and subcomponent name is not valid.
"""
if component not in Component:
raise ValueError("Given component name not found")
if subcomponent not in SubComponent:
raise ValueError("Given subcomponent name not found")
if (subcomponent != SubComponent.UNSPECIFIED and
subcomponent.component != component):
raise ValueError("component and subcomponent name don't match")
def report_error(error_data: converter_error_data_pb2.ConverterErrorData):
# Always overwrites the component information, but only overwrites the
# subcomponent if it is not available.
error_data.component = component.value
if not error_data.subcomponent:
error_data.subcomponent = subcomponent.name
tflite_metrics = metrics.TFLiteConverterMetrics()
tflite_metrics.set_converter_error(error_data)
def report_error_message(error_message: Text):
error_data = converter_error_data_pb2.ConverterErrorData()
error_data.error_message = error_message
report_error(error_data)
def actual_decorator(func):
@functools.wraps(func)
def wrapper(*args, **kwargs):
try:
return func(*args, **kwargs)
except ConverterError as converter_error:
if converter_error.errors:
for error_data in converter_error.errors:
report_error(error_data)
else:
report_error_message(str(converter_error))
raise converter_error from None # Re-throws the exception.
except Exception as error:
report_error_message(str(error))
raise error from None # Re-throws the exception.
return wrapper
return actual_decorator
tensorflow/lite/python/convert_saved_model.py
+186
View File
@@ -0,0 +1,186 @@
# Copyright 2018 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.
# ==============================================================================
"""Functions to convert SavedModel to frozen GraphDefs."""
from tensorflow.lite.python import util
from tensorflow.lite.python.convert_phase import Component
from tensorflow.lite.python.convert_phase import convert_phase
from tensorflow.lite.python.convert_phase import SubComponent
from tensorflow.python.client import session
from tensorflow.python.framework import ops
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.saved_model import constants
from tensorflow.python.saved_model import loader
def get_meta_graph_def(saved_model_dir, tag_set):
"""Validate saved_model and extract MetaGraphDef.
Args:
saved_model_dir: saved_model path to convert.
tag_set: Set of tag(s) of the MetaGraphDef to load.
Returns:
The meta_graph_def used for tflite conversion.
Raises:
ValueError: No valid MetaGraphDef for given tag_set.
"""
with session.Session(graph=ops.Graph()) as sess:
return loader.load(sess, tag_set, saved_model_dir)
def get_signature_def(meta_graph, signature_key):
"""Get the signature def from meta_graph with given signature_key.
Args:
meta_graph: meta_graph_def.
signature_key: signature_def in the meta_graph_def.
Returns:
The signature_def used for tflite conversion.
Raises:
ValueError: Given signature_key is not valid for this meta_graph.
"""
signature_def_map = meta_graph.signature_def
signature_def_keys = set(signature_def_map.keys())
logging.info(
"The given SavedModel MetaGraphDef contains SignatureDefs with the "
"following keys: %s", signature_def_keys)
if signature_key not in signature_def_keys:
raise ValueError("No '{}' in the SavedModel\'s SignatureDefs. Possible "
"values are '{}'.".format(signature_key,
",".join(signature_def_keys)))
return signature_def_map[signature_key]
def get_inputs_outputs(signature_def):
"""Get inputs and outputs from SignatureDef.
Args:
signature_def: SignatureDef in the meta_graph_def for conversion.
Returns:
The inputs and outputs in the graph for conversion.
"""
inputs_tensor_info = signature_def.inputs
outputs_tensor_info = signature_def.outputs
def gather_names(tensor_info):
return [tensor_info[key].name for key in tensor_info]
inputs = gather_names(inputs_tensor_info)
outputs = gather_names(outputs_tensor_info)
return inputs, outputs
def _get_tensors(graph, signature_def_tensor_names=None,
user_tensor_names=None):
"""Gets the tensors associated with the tensor names.
Either signature_def_tensor_names or user_tensor_names should be provided. If
the user provides tensors, the tensors associated with the user provided
tensor names are provided. Otherwise, the tensors associated with the names in
the SignatureDef are provided.
Args:
graph: GraphDef representing graph.
signature_def_tensor_names: Tensor names stored in either the inputs or
outputs of a SignatureDef. (default None)
user_tensor_names: Tensor names provided by the user. (default None)
Returns:
List of tensors.
Raises:
ValueError:
signature_def_tensors and user_tensor_names are undefined or empty.
user_tensor_names are not valid.
"""
tensors = []
if user_tensor_names:
# Sort the tensor names.
user_tensor_names = sorted(user_tensor_names)
tensors = util.get_tensors_from_tensor_names(graph, user_tensor_names)
elif signature_def_tensor_names:
tensors = [
graph.get_tensor_by_name(name)
for name in sorted(signature_def_tensor_names)
]
else:
# Throw ValueError if signature_def_tensors and user_tensor_names are both
# either undefined or empty.
raise ValueError(
"Specify either signature_def_tensor_names or user_tensor_names")
return tensors
@convert_phase(Component.PREPARE_TF_MODEL, SubComponent.FREEZE_SAVED_MODEL)
def freeze_saved_model(saved_model_dir, input_arrays, input_shapes,
output_arrays, tag_set, signature_key):
"""Converts a SavedModel to a frozen graph.
Args:
saved_model_dir: SavedModel directory to convert.
input_arrays: List of input tensors to freeze graph with. Uses input arrays
from SignatureDef when none are provided.
input_shapes: Dict of strings representing input tensor names to list of
integers representing input shapes (e.g., {"foo": : [1, 16, 16, 3]}).
Automatically determined when input shapes is None (e.g., {"foo" : None}).
output_arrays: List of output tensors to freeze graph with. Uses output
arrays from SignatureDef when none are provided.
tag_set: Set of tags identifying the MetaGraphDef within the SavedModel to
analyze. All tags in the tag set must be present.
signature_key: Key identifying SignatureDef containing inputs and outputs.
Returns:
frozen_graph_def: Frozen GraphDef.
in_tensors: List of input tensors for the graph.
out_tensors: List of output tensors for the graph.
graph: `Graph` object.
Raises:
ValueError:
SavedModel doesn't contain a MetaGraphDef identified by tag_set.
signature_key is not in the MetaGraphDef.
assets/ directory is in the MetaGraphDef.
input_shapes does not match the length of input_arrays.
input_arrays or output_arrays are not valid.
"""
# Read SignatureDef.
meta_graph = get_meta_graph_def(saved_model_dir, tag_set)
signature_def = get_signature_def(meta_graph, signature_key)
inputs, outputs = get_inputs_outputs(signature_def)
# Check SavedModel for assets directory.
collection_def = meta_graph.collection_def
if constants.ASSETS_KEY in collection_def:
raise ValueError("SavedModels with assets/ directory are not supported.")
graph = ops.Graph()
with session.Session(graph=graph) as sess:
loader.load(sess, meta_graph.meta_info_def.tags, saved_model_dir)
# Gets input and output tensors.
# TODO(zhixianyan): Use TFLite supported Op list to filter outputs.
in_tensors = _get_tensors(graph, inputs, input_arrays)
out_tensors = _get_tensors(graph, outputs, output_arrays)
util.set_tensor_shapes(in_tensors, input_shapes)
frozen_graph_def = util.freeze_graph(sess, in_tensors, out_tensors)
return frozen_graph_def, in_tensors, out_tensors, sess.graph
tensorflow/lite/python/convert_saved_model_test.py
+249
View File
@@ -0,0 +1,249 @@
# Copyright 2018 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.
# ==============================================================================
"""TFLite SavedModel conversion test cases.
- Tests converting simple SavedModel graph to TFLite FlatBuffer.
- Tests converting simple SavedModel graph to frozen graph.
- Tests converting MNIST SavedModel to TFLite FlatBuffer.
"""
import os
from tensorflow.lite.python import convert_saved_model
from tensorflow.python.client import session
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.platform import test
from tensorflow.python.saved_model import saved_model
from tensorflow.python.saved_model import signature_constants
from tensorflow.python.saved_model import tag_constants
class FreezeSavedModelTest(test_util.TensorFlowTestCase):
def _createSimpleSavedModel(self, shape):
"""Create a simple SavedModel on the fly."""
saved_model_dir = os.path.join(self.get_temp_dir(), "simple_savedmodel")
with session.Session() as sess:
in_tensor = array_ops.placeholder(shape=shape, dtype=dtypes.float32)
out_tensor = in_tensor + in_tensor
inputs = {"x": in_tensor}
outputs = {"y": out_tensor}
saved_model.simple_save(sess, saved_model_dir, inputs, outputs)
return saved_model_dir
def _createSavedModelTwoInputArrays(self, shape):
"""Create a simple SavedModel."""
saved_model_dir = os.path.join(self.get_temp_dir(), "simple_savedmodel")
with session.Session() as sess:
in_tensor_1 = array_ops.placeholder(
shape=shape, dtype=dtypes.float32, name="inputB")
in_tensor_2 = array_ops.placeholder(
shape=shape, dtype=dtypes.float32, name="inputA")
out_tensor = in_tensor_1 + in_tensor_2
inputs = {"x": in_tensor_1, "y": in_tensor_2}
outputs = {"z": out_tensor}
saved_model.simple_save(sess, saved_model_dir, inputs, outputs)
return saved_model_dir
def _getArrayNames(self, tensors):
return [tensor.name for tensor in tensors]
def _getArrayShapes(self, tensors):
dims = []
for tensor in tensors:
dim_tensor = []
for dim in tensor.shape:
if isinstance(dim, tensor_shape.Dimension):
dim_tensor.append(dim.value)
else:
dim_tensor.append(dim)
dims.append(dim_tensor)
return dims
def _convertSavedModel(self,
saved_model_dir,
input_arrays=None,
input_shapes=None,
output_arrays=None,
tag_set=None,
signature_key=None):
if tag_set is None:
tag_set = set([tag_constants.SERVING])
if signature_key is None:
signature_key = signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY
graph_def, in_tensors, out_tensors, _ = (
convert_saved_model.freeze_saved_model(
saved_model_dir=saved_model_dir,
input_arrays=input_arrays,
input_shapes=input_shapes,
output_arrays=output_arrays,
tag_set=tag_set,
signature_key=signature_key))
return graph_def, in_tensors, out_tensors
def testSimpleSavedModel(self):
"""Test a SavedModel."""
saved_model_dir = self._createSimpleSavedModel(shape=[1, 16, 16, 3])
_, in_tensors, out_tensors = self._convertSavedModel(saved_model_dir)
self.assertEqual(self._getArrayNames(out_tensors), ["add:0"])
self.assertEqual(self._getArrayNames(in_tensors), ["Placeholder:0"])
self.assertEqual(self._getArrayShapes(in_tensors), [[1, 16, 16, 3]])
def testSimpleSavedModelWithNoneBatchSizeInShape(self):
"""Test a SavedModel with None in input tensor's shape."""
saved_model_dir = self._createSimpleSavedModel(shape=[None, 16, 16, 3])
_, in_tensors, out_tensors = self._convertSavedModel(saved_model_dir)
self.assertEqual(self._getArrayNames(out_tensors), ["add:0"])
self.assertEqual(self._getArrayNames(in_tensors), ["Placeholder:0"])
self.assertEqual(self._getArrayShapes(in_tensors), [[None, 16, 16, 3]])
def testSimpleSavedModelWithInvalidSignatureKey(self):
"""Test a SavedModel that fails due to an invalid signature_key."""
saved_model_dir = self._createSimpleSavedModel(shape=[1, 16, 16, 3])
with self.assertRaises(ValueError) as error:
self._convertSavedModel(saved_model_dir, signature_key="invalid-key")
self.assertEqual(
"No 'invalid-key' in the SavedModel's SignatureDefs. "
"Possible values are 'serving_default'.", str(error.exception))
def testSimpleSavedModelWithInvalidOutputArray(self):
"""Test a SavedModel that fails due to invalid output arrays."""
saved_model_dir = self._createSimpleSavedModel(shape=[1, 16, 16, 3])
with self.assertRaises(ValueError) as error:
self._convertSavedModel(saved_model_dir, output_arrays=["invalid-output"])
self.assertEqual("Invalid tensors 'invalid-output' were found.",
str(error.exception))
def testSimpleSavedModelWithWrongInputArrays(self):
"""Test a SavedModel that fails due to invalid input arrays."""
saved_model_dir = self._createSimpleSavedModel(shape=[1, 16, 16, 3])
# Check invalid input_arrays.
with self.assertRaises(ValueError) as error:
self._convertSavedModel(saved_model_dir, input_arrays=["invalid-input"])
self.assertEqual("Invalid tensors 'invalid-input' were found.",
str(error.exception))
# Check valid and invalid input_arrays.
with self.assertRaises(ValueError) as error:
self._convertSavedModel(
saved_model_dir, input_arrays=["Placeholder", "invalid-input"])
self.assertEqual("Invalid tensors 'invalid-input' were found.",
str(error.exception))
def testSimpleSavedModelWithCorrectArrays(self):
"""Test a SavedModel with correct input_arrays and output_arrays."""
saved_model_dir = self._createSimpleSavedModel(shape=[None, 16, 16, 3])
_, in_tensors, out_tensors = self._convertSavedModel(
saved_model_dir=saved_model_dir,
input_arrays=["Placeholder"],
output_arrays=["add"])
self.assertEqual(self._getArrayNames(out_tensors), ["add:0"])
self.assertEqual(self._getArrayNames(in_tensors), ["Placeholder:0"])
self.assertEqual(self._getArrayShapes(in_tensors), [[None, 16, 16, 3]])
def testSimpleSavedModelWithCorrectInputArrays(self):
"""Test a SavedModel with correct input_arrays and input_shapes."""
saved_model_dir = self._createSimpleSavedModel(shape=[1, 16, 16, 3])
_, in_tensors, out_tensors = self._convertSavedModel(
saved_model_dir=saved_model_dir,
input_arrays=["Placeholder"],
input_shapes={"Placeholder": [1, 16, 16, 3]})
self.assertEqual(self._getArrayNames(out_tensors), ["add:0"])
self.assertEqual(self._getArrayNames(in_tensors), ["Placeholder:0"])
self.assertEqual(self._getArrayShapes(in_tensors), [[1, 16, 16, 3]])
def testTwoInputArrays(self):
"""Test a simple SavedModel."""
saved_model_dir = self._createSavedModelTwoInputArrays(shape=[1, 16, 16, 3])
_, in_tensors, out_tensors = self._convertSavedModel(
saved_model_dir=saved_model_dir, input_arrays=["inputB", "inputA"])
self.assertEqual(self._getArrayNames(out_tensors), ["add:0"])
self.assertEqual(self._getArrayNames(in_tensors), ["inputA:0", "inputB:0"])
self.assertEqual(
self._getArrayShapes(in_tensors), [[1, 16, 16, 3], [1, 16, 16, 3]])
def testSubsetInputArrays(self):
"""Test a SavedModel with a subset of the input array names of the model."""
saved_model_dir = self._createSavedModelTwoInputArrays(shape=[1, 16, 16, 3])
# Check case where input shape is given.
_, in_tensors, out_tensors = self._convertSavedModel(
saved_model_dir=saved_model_dir,
input_arrays=["inputA"],
input_shapes={"inputA": [1, 16, 16, 3]})
self.assertEqual(self._getArrayNames(out_tensors), ["add:0"])
self.assertEqual(self._getArrayNames(in_tensors), ["inputA:0"])
self.assertEqual(self._getArrayShapes(in_tensors), [[1, 16, 16, 3]])
# Check case where input shape is None.
_, in_tensors, out_tensors = self._convertSavedModel(
saved_model_dir=saved_model_dir, input_arrays=["inputA"])
self.assertEqual(self._getArrayNames(out_tensors), ["add:0"])
self.assertEqual(self._getArrayNames(in_tensors), ["inputA:0"])
self.assertEqual(self._getArrayShapes(in_tensors), [[1, 16, 16, 3]])
def testMultipleMetaGraphDef(self):
"""Test saved model with multiple MetaGraphDefs."""
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
_, in_tensors, out_tensors = self._convertSavedModel(
saved_model_dir=saved_model_dir,
tag_set=set([saved_model.tag_constants.SERVING, "additional_test_tag"]))
self.assertEqual(self._getArrayNames(out_tensors), ["add:0"])
self.assertEqual(self._getArrayNames(in_tensors), ["Placeholder:0"])
self.assertEqual(self._getArrayShapes(in_tensors), [[1, 28, 28]])
if __name__ == "__main__":
test.main()
tensorflow/lite/python/convert_test.py
+495
View File
@@ -0,0 +1,495 @@
# 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()
tensorflow/lite/python/dummy_pybind.cc
+17
View File
@@ -0,0 +1,17 @@
/* Copyright 2025 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.
==============================================================================*/
#include <pybind11/pybind11.h>
PYBIND11_MODULE(_dummy_pybind, m) { m.doc() = "Dummy pybind module."; }
tensorflow/lite/python/interpreter.py
+1064
View File
File diff suppressed because it is too large Load Diff
tensorflow/lite/python/interpreter_test.py
+686
View File
@@ -0,0 +1,686 @@
# Copyright 2018 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."""
import ctypes
import io
import pathlib
import sys
from unittest import mock
import numpy as np
import tensorflow as tf
# Force loaded shared object symbols to be globally visible. This is needed so
# that the interpreter_wrapper, in one .so file, can see the test_registerer,
# in a different .so file. Note that this may already be set by default.
# pylint: disable=g-import-not-at-top
if hasattr(sys, 'setdlopenflags') and hasattr(sys, 'getdlopenflags'):
sys.setdlopenflags(sys.getdlopenflags() | ctypes.RTLD_GLOBAL)
from tensorflow.lite.python import interpreter as interpreter_wrapper
from tensorflow.lite.python import lite
from tensorflow.lite.python.metrics import metrics
from tensorflow.lite.python.testdata import _pywrap_test_registerer as test_registerer
from tensorflow.python.framework import test_util
from tensorflow.python.platform import resource_loader
from tensorflow.python.platform import test
# pylint: enable=g-import-not-at-top
class InterpreterCustomOpsTest(test_util.TensorFlowTestCase):
def testRegistererByName(self):
interpreter = interpreter_wrapper.InterpreterWithCustomOps(
model_path=resource_loader.get_path_to_datafile(
'testdata/permute_float.tflite'),
custom_op_registerers=['TF_TestRegisterer'])
self.assertTrue(interpreter._safe_to_run())
self.assertEqual(test_registerer.get_num_test_registerer_calls(), 1)
def testRegistererByFunc(self):
interpreter = interpreter_wrapper.InterpreterWithCustomOps(
model_path=resource_loader.get_path_to_datafile(
'testdata/permute_float.tflite'),
custom_op_registerers=[test_registerer.TF_TestRegisterer])
self.assertTrue(interpreter._safe_to_run())
self.assertEqual(test_registerer.get_num_test_registerer_calls(), 1)
def testRegistererFailure(self):
bogus_name = 'CompletelyBogusRegistererName'
with self.assertRaisesRegex(
ValueError, 'Looking up symbol \'' + bogus_name + '\' failed'):
interpreter_wrapper.InterpreterWithCustomOps(
model_path=resource_loader.get_path_to_datafile(
'testdata/permute_float.tflite'),
custom_op_registerers=[bogus_name])
# Register GenAI Ops is only supported when using LiteRT wheel.
def testRegisterGenAIOpsFailure(self):
genai_ops_name = 'pywrap_genai_ops.GenAIOpsRegisterer'
with self.assertRaisesRegex(
ValueError,
"Loading library 'pywrap_genai_ops.so' failed with error"
" 'pywrap_genai_ops.so: cannot open shared object file: No such file or"
" directory'",
):
interpreter_wrapper.InterpreterWithCustomOps(
model_path=resource_loader.get_path_to_datafile(
'testdata/permute_float.tflite'
),
custom_op_registerers=[genai_ops_name],
)
def testNoCustomOps(self):
interpreter = interpreter_wrapper.InterpreterWithCustomOps(
model_path=resource_loader.get_path_to_datafile(
'testdata/permute_float.tflite'))
self.assertTrue(interpreter._safe_to_run())
class InterpreterTest(test_util.TensorFlowTestCase):
def assertQuantizationParamsEqual(self, scales, zero_points,
quantized_dimension, params):
self.assertAllEqual(scales, params['scales'])
self.assertAllEqual(zero_points, params['zero_points'])
self.assertEqual(quantized_dimension, params['quantized_dimension'])
def testPathLikeModel(self):
interpreter = interpreter_wrapper.Interpreter(
model_path=pathlib.Path(
resource_loader.get_path_to_datafile(
'testdata/permute_float.tflite'
)
),
)
interpreter.allocate_tensors()
def testThreads_NegativeValue(self):
with self.assertRaisesRegex(ValueError, 'num_threads should >= 1'):
interpreter_wrapper.Interpreter(
model_path=resource_loader.get_path_to_datafile(
'testdata/permute_float.tflite'),
num_threads=-1)
def testThreads_WrongType(self):
with self.assertRaisesRegex(ValueError,
'type of num_threads should be int'):
interpreter_wrapper.Interpreter(
model_path=resource_loader.get_path_to_datafile(
'testdata/permute_float.tflite'),
num_threads=4.2)
def testNotSupportedOpResolverTypes(self):
with self.assertRaisesRegex(
ValueError, 'Unrecognized passed in op resolver type: test'):
interpreter_wrapper.Interpreter(
model_path=resource_loader.get_path_to_datafile(
'testdata/permute_float.tflite'),
experimental_op_resolver_type='test')
def testFloatWithDifferentOpResolverTypes(self):
op_resolver_types = [
interpreter_wrapper.OpResolverType.BUILTIN,
interpreter_wrapper.OpResolverType.BUILTIN_REF,
interpreter_wrapper.OpResolverType.BUILTIN_WITHOUT_DEFAULT_DELEGATES
]
for op_resolver_type in op_resolver_types:
interpreter = interpreter_wrapper.Interpreter(
model_path=resource_loader.get_path_to_datafile(
'testdata/permute_float.tflite'),
experimental_op_resolver_type=op_resolver_type)
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, 4] == input_details[0]['shape']).all())
self.assertEqual((0.0, 0), input_details[0]['quantization'])
self.assertQuantizationParamsEqual(
[], [], 0, input_details[0]['quantization_parameters'])
output_details = interpreter.get_output_details()
self.assertEqual(1, len(output_details))
self.assertEqual('output', output_details[0]['name'])
self.assertEqual(np.float32, output_details[0]['dtype'])
self.assertTrue(([1, 4] == output_details[0]['shape']).all())
self.assertEqual((0.0, 0), output_details[0]['quantization'])
self.assertQuantizationParamsEqual(
[], [], 0, output_details[0]['quantization_parameters'])
test_input = np.array([[1.0, 2.0, 3.0, 4.0]], dtype=np.float32)
expected_output = np.array([[4.0, 3.0, 2.0, 1.0]], dtype=np.float32)
interpreter.set_tensor(input_details[0]['index'], test_input)
interpreter.invoke()
output_data = interpreter.get_tensor(output_details[0]['index'])
self.assertTrue((expected_output == output_data).all())
def testFloatWithTwoThreads(self):
interpreter = interpreter_wrapper.Interpreter(
model_path=resource_loader.get_path_to_datafile(
'testdata/permute_float.tflite'),
num_threads=2)
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
test_input = np.array([[1.0, 2.0, 3.0, 4.0]], dtype=np.float32)
expected_output = np.array([[4.0, 3.0, 2.0, 1.0]], dtype=np.float32)
interpreter.set_tensor(input_details[0]['index'], test_input)
interpreter.invoke()
output_details = interpreter.get_output_details()
output_data = interpreter.get_tensor(output_details[0]['index'])
self.assertTrue((expected_output == output_data).all())
def testUint8(self):
model_path = resource_loader.get_path_to_datafile(
'testdata/permute_uint8.tflite')
with io.open(model_path, 'rb') as model_file:
data = model_file.read()
interpreter = interpreter_wrapper.Interpreter(model_content=data)
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.uint8, input_details[0]['dtype'])
self.assertTrue(([1, 4] == input_details[0]['shape']).all())
self.assertEqual((1.0, 0), input_details[0]['quantization'])
self.assertQuantizationParamsEqual(
[1.0], [0], 0, input_details[0]['quantization_parameters'])
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, 4] == output_details[0]['shape']).all())
self.assertEqual((1.0, 0), output_details[0]['quantization'])
self.assertQuantizationParamsEqual(
[1.0], [0], 0, output_details[0]['quantization_parameters'])
test_input = np.array([[1, 2, 3, 4]], dtype=np.uint8)
expected_output = np.array([[4, 3, 2, 1]], dtype=np.uint8)
interpreter.resize_tensor_input(input_details[0]['index'], test_input.shape)
interpreter.allocate_tensors()
interpreter.set_tensor(input_details[0]['index'], test_input)
interpreter.invoke()
output_data = interpreter.get_tensor(output_details[0]['index'])
self.assertTrue((expected_output == output_data).all())
def testString(self):
interpreter = interpreter_wrapper.Interpreter(
model_path=resource_loader.get_path_to_datafile(
'testdata/gather_string.tflite'))
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
self.assertEqual(2, len(input_details))
self.assertEqual('input', input_details[0]['name'])
self.assertEqual(np.bytes_, input_details[0]['dtype'])
self.assertTrue(([10] == input_details[0]['shape']).all())
self.assertEqual((0.0, 0), input_details[0]['quantization'])
self.assertQuantizationParamsEqual(
[], [], 0, input_details[0]['quantization_parameters'])
self.assertEqual('indices', input_details[1]['name'])
self.assertEqual(np.int64, input_details[1]['dtype'])
self.assertTrue(([3] == input_details[1]['shape']).all())
self.assertEqual((0.0, 0), input_details[1]['quantization'])
self.assertQuantizationParamsEqual(
[], [], 0, input_details[1]['quantization_parameters'])
output_details = interpreter.get_output_details()
self.assertEqual(1, len(output_details))
self.assertEqual('output', output_details[0]['name'])
self.assertEqual(np.bytes_, output_details[0]['dtype'])
self.assertTrue(([3] == output_details[0]['shape']).all())
self.assertEqual((0.0, 0), output_details[0]['quantization'])
self.assertQuantizationParamsEqual(
[], [], 0, output_details[0]['quantization_parameters'])
test_input = np.array([1, 2, 3], dtype=np.int64)
interpreter.set_tensor(input_details[1]['index'], test_input)
test_input = np.array(['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j'])
expected_output = np.array([b'b', b'c', b'd'])
interpreter.set_tensor(input_details[0]['index'], test_input)
interpreter.invoke()
output_data = interpreter.get_tensor(output_details[0]['index'])
self.assertTrue((expected_output == output_data).all())
def testStringZeroDim(self):
data = b'abcd' + bytes(16)
interpreter = interpreter_wrapper.Interpreter(
model_path=resource_loader.get_path_to_datafile(
'testdata/gather_string_0d.tflite'))
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
interpreter.set_tensor(input_details[0]['index'], np.array(data))
test_input_tensor = interpreter.get_tensor(input_details[0]['index'])
self.assertEqual(len(data), len(test_input_tensor.item(0)))
def testPerChannelParams(self):
interpreter = interpreter_wrapper.Interpreter(
model_path=resource_loader.get_path_to_datafile('testdata/pc_conv.bin'))
interpreter.allocate_tensors()
# Tensor index 1 is the weight.
weight_details = interpreter.get_tensor_details()[1]
qparams = weight_details['quantization_parameters']
# Ensure that we retrieve per channel quantization params correctly.
self.assertEqual(len(qparams['scales']), 128)
def testDenseTensorAccess(self):
interpreter = interpreter_wrapper.Interpreter(
model_path=resource_loader.get_path_to_datafile('testdata/pc_conv.bin'))
interpreter.allocate_tensors()
weight_details = interpreter.get_tensor_details()[1]
s_params = weight_details['sparsity_parameters']
self.assertEqual(s_params, {})
def testSparseTensorAccess(self):
interpreter = interpreter_wrapper.InterpreterWithCustomOps(
model_path=resource_loader.get_path_to_datafile(
'../testdata/sparse_tensor.bin'),
custom_op_registerers=['TF_TestRegisterer'])
interpreter.allocate_tensors()
# Tensor at index 0 is sparse.
compressed_buffer = interpreter.get_tensor(0)
# Ensure that the buffer is of correct size and value.
self.assertEqual(len(compressed_buffer), 12)
sparse_value = [1, 0, 0, 4, 2, 3, 0, 0, 5, 0, 0, 6]
self.assertAllEqual(compressed_buffer, sparse_value)
tensor_details = interpreter.get_tensor_details()[0]
s_params = tensor_details['sparsity_parameters']
# Ensure sparsity parameter returned is correct
self.assertAllEqual(s_params['traversal_order'], [0, 1, 2, 3])
self.assertAllEqual(s_params['block_map'], [0, 1])
dense_dim_metadata = {'format': 0, 'dense_size': 2}
self.assertAllEqual(s_params['dim_metadata'][0], dense_dim_metadata)
self.assertAllEqual(s_params['dim_metadata'][2], dense_dim_metadata)
self.assertAllEqual(s_params['dim_metadata'][3], dense_dim_metadata)
self.assertEqual(s_params['dim_metadata'][1]['format'], 1)
self.assertAllEqual(s_params['dim_metadata'][1]['array_segments'],
[0, 2, 3])
self.assertAllEqual(s_params['dim_metadata'][1]['array_indices'], [0, 1, 1])
@mock.patch.object(metrics.TFLiteMetrics,
'increase_counter_interpreter_creation')
def testCreationCounter(self, increase_call):
interpreter_wrapper.Interpreter(
model_path=resource_loader.get_path_to_datafile(
'testdata/permute_float.tflite'))
increase_call.assert_called_once()
class InterpreterTestErrorPropagation(test_util.TensorFlowTestCase):
# Model must have at least 7 bytes to hold model identifier
def testTooShortModelContent(self):
with self.assertRaisesRegex(ValueError,
'The model is not a valid Flatbuffer buffer'):
interpreter_wrapper.Interpreter(model_content=b'short')
def testInvalidModelContent(self):
with self.assertRaisesRegex(ValueError,
'The model is not a valid Flatbuffer buffer'):
interpreter_wrapper.Interpreter(model_content=b'wrong_identifier')
def testInvalidModelFile(self):
with self.assertRaisesRegex(ValueError,
'Could not open \'totally_invalid_file_name\''):
interpreter_wrapper.Interpreter(model_path='totally_invalid_file_name')
def testInvokeBeforeReady(self):
interpreter = interpreter_wrapper.Interpreter(
model_path=resource_loader.get_path_to_datafile(
'testdata/permute_float.tflite'))
with self.assertRaisesRegex(RuntimeError,
'Invoke called on model that is not ready'):
interpreter.invoke()
def testInvalidModelFileContent(self):
with self.assertRaisesRegex(
ValueError, '`model_path` or `model_content` must be specified.'):
interpreter_wrapper.Interpreter(model_path=None, model_content=None)
def testInvalidIndex(self):
interpreter = interpreter_wrapper.Interpreter(
model_path=resource_loader.get_path_to_datafile(
'testdata/permute_float.tflite'))
interpreter.allocate_tensors()
# Invalid tensor index passed.
with self.assertRaisesRegex(
ValueError, 'Invalid tensor index 4 exceeds max tensor index 3'
):
interpreter._get_tensor_details(4, 0)
with self.assertRaisesRegex(ValueError, 'Invalid node index'):
interpreter._get_op_details(4)
def testEmptyInputTensor(self):
class TestModel(tf.keras.models.Model):
@tf.function(
input_signature=[tf.TensorSpec(shape=[None], dtype=tf.float32)])
def TestSum(self, x):
return tf.raw_ops.Sum(input=x, axis=[0])
test_model = TestModel()
converter = lite.TFLiteConverterV2.from_concrete_functions([
test_model.TestSum.get_concrete_function(
tf.TensorSpec([None], tf.float32))
], test_model)
model = converter.convert()
interpreter = lite.Interpreter(model_content=model)
# Make sure that passing empty tensor doesn't cause any errors.
interpreter.get_signature_runner()(x=tf.zeros([0], tf.float32))
class InterpreterTensorAccessorTest(test_util.TensorFlowTestCase):
def setUp(self):
super(InterpreterTensorAccessorTest, self).setUp()
self.interpreter = interpreter_wrapper.Interpreter(
model_path=resource_loader.get_path_to_datafile(
'testdata/permute_float.tflite'))
self.interpreter.allocate_tensors()
self.input0 = self.interpreter.get_input_details()[0]['index']
self.initial_data = np.array([[-1., -2., -3., -4.]], np.float32)
def testTensorAccessor(self):
"""Check that tensor returns a reference."""
array_ref = self.interpreter.tensor(self.input0)
np.copyto(array_ref(), self.initial_data)
self.assertAllEqual(array_ref(), self.initial_data)
self.assertAllEqual(
self.interpreter.get_tensor(self.input0), self.initial_data)
def testGetTensorAccessor(self):
"""Check that get_tensor returns a copy."""
self.interpreter.set_tensor(self.input0, self.initial_data)
array_initial_copy = self.interpreter.get_tensor(self.input0)
new_value = np.add(1., array_initial_copy)
self.interpreter.set_tensor(self.input0, new_value)
self.assertAllEqual(array_initial_copy, self.initial_data)
self.assertAllEqual(self.interpreter.get_tensor(self.input0), new_value)
def testBase(self):
self.assertTrue(self.interpreter._safe_to_run())
_ = self.interpreter.tensor(self.input0)
self.assertTrue(self.interpreter._safe_to_run())
in0 = self.interpreter.tensor(self.input0)()
self.assertFalse(self.interpreter._safe_to_run())
in0b = self.interpreter.tensor(self.input0)()
self.assertFalse(self.interpreter._safe_to_run())
# Now get rid of the buffers so that we can evaluate.
del in0
del in0b
self.assertTrue(self.interpreter._safe_to_run())
def testBaseProtectsFunctions(self):
in0 = self.interpreter.tensor(self.input0)()
# Make sure we get an exception if we try to run an unsafe operation
with self.assertRaisesRegex(RuntimeError, 'There is at least 1 reference'):
_ = self.interpreter.allocate_tensors()
# Make sure we get an exception if we try to run an unsafe operation
with self.assertRaisesRegex(RuntimeError, 'There is at least 1 reference'):
_ = self.interpreter.invoke() # pylint: disable=assignment-from-no-return
# Now test that we can run
del in0 # this is our only buffer reference, so now it is safe to change
in0safe = self.interpreter.tensor(self.input0)
_ = self.interpreter.allocate_tensors()
del in0safe # make sure in0Safe is held but lint doesn't complain
class InterpreterNodeAccessTest(test_util.TensorFlowTestCase):
def setUp(self):
super().setUp()
self.interpreter = interpreter_wrapper.Interpreter(
model_path=resource_loader.get_path_to_datafile(
'testdata/permute_float.tflite'
)
)
self.interpreter.allocate_tensors()
self.input0 = self.interpreter.get_input_details()[0]['index']
self.initial_data = np.array([[-1.0, -2.0, -3.0, -4.0]], np.float32)
def testValidNode(self):
"""Check that tensor returns a reference."""
ops_details = self.interpreter._get_ops_details()
self.assertEqual(ops_details[0]['index'], 0)
self.assertEqual(ops_details[0]['op_name'], 'FULLY_CONNECTED')
self.assertAllEqual(ops_details[0]['inputs'], [0, 1, -1])
self.assertAllEqual(ops_details[0]['outputs'], [2])
self.assertAllEqual(
ops_details[0]['operand_types'], [np.float32, np.float32]
)
self.assertAllEqual(ops_details[0]['result_types'], [np.float32])
def testInvalidNode(self):
with self.assertRaisesRegex(ValueError, 'Invalid node index'):
self.interpreter._get_op_details(4)
class InterpreterDelegateTest(test_util.TensorFlowTestCase):
def setUp(self):
super(InterpreterDelegateTest, self).setUp()
self._delegate_file = resource_loader.get_path_to_datafile(
'testdata/test_delegate.so')
self._model_file = resource_loader.get_path_to_datafile(
'testdata/permute_float.tflite')
# Load the library to reset the counters.
library = ctypes.pydll.LoadLibrary(self._delegate_file)
library.initialize_counters()
def _TestInterpreter(self, model_path, options=None):
"""Test wrapper function that creates an interpreter with the delegate."""
delegate = interpreter_wrapper.load_delegate(self._delegate_file, options)
return interpreter_wrapper.Interpreter(
model_path=model_path, experimental_delegates=[delegate])
def testDelegate(self):
"""Tests the delegate creation and destruction."""
interpreter = self._TestInterpreter(model_path=self._model_file)
lib = interpreter._delegates[0]._library
self.assertEqual(lib.get_num_delegates_created(), 1)
self.assertEqual(lib.get_num_delegates_destroyed(), 0)
self.assertEqual(lib.get_num_delegates_invoked(), 1)
del interpreter
self.assertEqual(lib.get_num_delegates_created(), 1)
self.assertEqual(lib.get_num_delegates_destroyed(), 1)
self.assertEqual(lib.get_num_delegates_invoked(), 1)
def testMultipleInterpreters(self):
delegate = interpreter_wrapper.load_delegate(self._delegate_file)
lib = delegate._library
self.assertEqual(lib.get_num_delegates_created(), 1)
self.assertEqual(lib.get_num_delegates_destroyed(), 0)
self.assertEqual(lib.get_num_delegates_invoked(), 0)
interpreter_a = interpreter_wrapper.Interpreter(
model_path=self._model_file, experimental_delegates=[delegate])
self.assertEqual(lib.get_num_delegates_created(), 1)
self.assertEqual(lib.get_num_delegates_destroyed(), 0)
self.assertEqual(lib.get_num_delegates_invoked(), 1)
interpreter_b = interpreter_wrapper.Interpreter(
model_path=self._model_file, experimental_delegates=[delegate])
self.assertEqual(lib.get_num_delegates_created(), 1)
self.assertEqual(lib.get_num_delegates_destroyed(), 0)
self.assertEqual(lib.get_num_delegates_invoked(), 2)
del delegate
del interpreter_a
self.assertEqual(lib.get_num_delegates_created(), 1)
self.assertEqual(lib.get_num_delegates_destroyed(), 0)
self.assertEqual(lib.get_num_delegates_invoked(), 2)
del interpreter_b
self.assertEqual(lib.get_num_delegates_created(), 1)
self.assertEqual(lib.get_num_delegates_destroyed(), 1)
self.assertEqual(lib.get_num_delegates_invoked(), 2)
def testDestructionOrder(self):
"""Make sure internal _interpreter object is destroyed before delegate."""
self.skipTest('TODO(b/142136355): fix flakiness and re-enable')
# Track which order destructions were doned in
destructions = []
def register_destruction(x):
destructions.append(x if isinstance(x, str) else x.decode('utf-8'))
return 0
# Make a wrapper for the callback so we can send this to ctypes
delegate = interpreter_wrapper.load_delegate(self._delegate_file)
# Make an interpreter with the delegate
interpreter = interpreter_wrapper.Interpreter(
model_path=resource_loader.get_path_to_datafile(
'testdata/permute_float.tflite'),
experimental_delegates=[delegate])
class InterpreterDestroyCallback:
def __del__(self):
register_destruction('interpreter')
interpreter._interpreter.stuff = InterpreterDestroyCallback()
# Destroy both delegate and interpreter
library = delegate._library
prototype = ctypes.CFUNCTYPE(ctypes.c_int, (ctypes.c_char_p))
library.set_destroy_callback(prototype(register_destruction))
del delegate
del interpreter
library.set_destroy_callback(None)
# check the interpreter was destroyed before the delegate
self.assertEqual(destructions, ['interpreter', 'test_delegate'])
def testOptions(self):
delegate_a = interpreter_wrapper.load_delegate(self._delegate_file)
lib = delegate_a._library
self.assertEqual(lib.get_num_delegates_created(), 1)
self.assertEqual(lib.get_num_delegates_destroyed(), 0)
self.assertEqual(lib.get_num_delegates_invoked(), 0)
self.assertEqual(lib.get_options_counter(), 0)
delegate_b = interpreter_wrapper.load_delegate(
self._delegate_file, options={
'unused': False,
'options_counter': 2
})
lib = delegate_b._library
self.assertEqual(lib.get_num_delegates_created(), 2)
self.assertEqual(lib.get_num_delegates_destroyed(), 0)
self.assertEqual(lib.get_num_delegates_invoked(), 0)
self.assertEqual(lib.get_options_counter(), 2)
del delegate_a
del delegate_b
self.assertEqual(lib.get_num_delegates_created(), 2)
self.assertEqual(lib.get_num_delegates_destroyed(), 2)
self.assertEqual(lib.get_num_delegates_invoked(), 0)
self.assertEqual(lib.get_options_counter(), 2)
def testFail(self):
with self.assertRaisesRegex(
# Due to exception chaining in PY3, we can't be more specific here and
# check that the phrase 'Fail argument sent' is present.
ValueError, 'Failed to load delegate from'):
interpreter_wrapper.load_delegate(
self._delegate_file, options={'fail': 'fail'})
class InterpreterMultiSignatureTest(test_util.TensorFlowTestCase):
def setUp(self):
super(InterpreterMultiSignatureTest, self).setUp()
self._single_signature_file = resource_loader.get_path_to_datafile(
'testdata/permute_float.tflite'
)
self._double_signature_file = resource_loader.get_path_to_datafile(
'testdata/two_signatures.tflite'
)
def testNumSubgraphsSingleSignature(self):
single_signature_interpreter = interpreter_wrapper.Interpreter(
model_path=self._single_signature_file
)
self.assertEqual(single_signature_interpreter.num_subgraphs(), 1)
def testNumSubgraphsDoubleSignature(self):
double_signature_interpreter = interpreter_wrapper.Interpreter(
model_path=self._double_signature_file
)
self.assertEqual(double_signature_interpreter.num_subgraphs(), 2)
def testGetTensorDetailsSingleSignature(self):
single_signature_interpreter = interpreter_wrapper.Interpreter(
model_path=self._single_signature_file
)
tensor_details = single_signature_interpreter.get_tensor_details()
self.assertLen(tensor_details, 3)
self.assertEqual(tensor_details[0]['name'], 'input')
with self.assertRaisesRegex(ValueError, 'subgraph_index is out of range'):
single_signature_interpreter.get_tensor_details(subgraph_index=1)
with self.assertRaisesRegex(ValueError, 'subgraph_index is out of range'):
single_signature_interpreter.get_tensor_details(subgraph_index=-1)
def testGetTensorDetailsDoubleSignature(self):
double_signature_interpreter = interpreter_wrapper.Interpreter(
model_path=self._double_signature_file
)
subgraph0_tensor_details = double_signature_interpreter.get_tensor_details(
subgraph_index=0
)
self.assertLen(subgraph0_tensor_details, 3)
self.assertEqual(subgraph0_tensor_details[0]['name'], 'add_x:0')
subgraph1_tensor_details = double_signature_interpreter.get_tensor_details(
subgraph_index=1
)
self.assertLen(subgraph1_tensor_details, 3)
self.assertEqual(subgraph1_tensor_details[0]['name'], 'multiply_x:0')
with self.assertRaisesRegex(ValueError, 'subgraph_index is out of range'):
double_signature_interpreter.get_tensor_details(subgraph_index=3)
if __name__ == '__main__':
test.main()
tensorflow/lite/python/interpreter_wrapper/BUILD
+112
View File
@@ -0,0 +1,112 @@
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable", "pybind_extension")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
default_visibility = ["//visibility:public"],
licenses = ["notice"],
)
cc_library(
name = "numpy",
srcs = ["numpy.cc"],
hdrs = ["numpy.h"],
compatible_with = get_compatible_with_portable(),
deps = [
"//tensorflow/lite:string_util",
"//tensorflow/lite/core/c:c_api_types",
"//tensorflow/lite/core/c:common",
"//third_party/py/numpy:headers",
"@xla//third_party/python_runtime:headers", # buildcleaner: keep
],
)
cc_library(
name = "interpreter_wrapper_lib",
srcs = ["interpreter_wrapper.cc"],
hdrs = [
"interpreter_wrapper.h",
],
compatible_with = get_compatible_with_portable(),
deps = [
":numpy",
":python_error_reporter",
":python_utils",
"//tensorflow/lite:framework",
"//tensorflow/lite:shared_library",
"//tensorflow/lite:string_util",
"//tensorflow/lite:util",
"//tensorflow/lite/core:framework",
"//tensorflow/lite/core/api",
"//tensorflow/lite/core/c:common",
"//tensorflow/lite/core/kernels:builtin_ops",
"//tensorflow/lite/delegates/xnnpack:xnnpack_delegate",
"//tensorflow/lite/kernels:reference_ops",
"//tensorflow/lite/kernels/internal:compatibility",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:str_format",
"@xla//third_party/python_runtime:headers", # buildcleaner: keep
],
)
cc_library(
name = "python_error_reporter",
srcs = ["python_error_reporter.cc"],
hdrs = ["python_error_reporter.h"],
compatible_with = get_compatible_with_portable(),
deps = [
"//tensorflow/lite:stateful_error_reporter",
"@xla//third_party/python_runtime:headers", # buildcleaner: keep
],
)
cc_library(
name = "python_utils",
srcs = ["python_utils.cc"],
hdrs = ["python_utils.h"],
compatible_with = get_compatible_with_portable(),
deps = [
"@xla//third_party/python_runtime:headers", # buildcleaner: keep
],
)
config_setting(
name = "tflite_pip_with_flex",
define_values = {
"tflite_pip_with_flex": "true",
},
)
pybind_extension(
name = "_pywrap_tensorflow_interpreter_wrapper",
srcs = [
"interpreter_wrapper_pybind11.cc",
],
hdrs = ["interpreter_wrapper.h"],
additional_stubgen_deps = [
"//third_party/py/numpy:numpy",
],
common_lib_packages = [
"litert/python",
"tensorflow/lite/python",
],
compatible_with = get_compatible_with_portable(),
enable_stub_generation = True,
link_in_framework = True,
pytype_srcs = [
"_pywrap_tensorflow_interpreter_wrapper.pyi",
],
wrap_py_init = True,
deps = [
":interpreter_wrapper_lib",
"//tensorflow/lite:framework",
"//tensorflow/lite/core:framework_stable",
"//tensorflow/python/lib/core:pybind11_lib",
"@pybind11",
"@xla//third_party/python_runtime:headers",
] + select({
":tflite_pip_with_flex": ["//tensorflow/lite/delegates/flex:delegate"],
"//conditions:default": [],
}),
)
tensorflow/lite/python/interpreter_wrapper/_pywrap_tensorflow_interpreter_wrapper.pyi
+47
View File
@@ -0,0 +1,47 @@
# Copyright 2023 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.
# ==============================================================================
class InterpreterWrapper:
def __init__(self, *args, **kwargs) -> None: ...
def AllocateTensors(self, subgraph_index: int = ...) -> object: ...
def GetSignatureDefs(self) -> object: ...
def GetSubgraphIndexFromSignature(self, arg0: str) -> object: ...
def GetTensor(self, tensor_index: int, subgraph_index: int = ...) -> object: ...
def InputIndices(self) -> object: ...
def Invoke(self, subgraph_index: int = ...) -> object: ...
def ModifyGraphWithDelegate(self, arg0: int) -> object: ...
def NodeInputs(self, arg0: int) -> object: ...
def NodeName(self, arg0: int) -> str: ...
def NodeOutputs(self, arg0: int) -> object: ...
def NumNodes(self) -> int: ...
def NumSubgraphs(self) -> int: ...
def NumTensors(self, arg0: int) -> int: ...
def OutputIndices(self) -> object: ...
def ResetVariableTensors(self) -> object: ...
def ResizeInputTensor(self, i: int, value: object, strict: bool, subgraph_index: int = ...) -> object: ...
def SetNumThreads(self, arg0: int) -> object: ...
def SetTensor(self, i: int, value: object, subgraph_index: int = ...) -> object: ...
def TensorName(self, arg0: int, arg1: int) -> str: ...
def TensorQuantization(self, arg0: int, arg1: int) -> object: ...
def TensorQuantizationParameters(self, arg0: int, arg1: int) -> object: ...
def TensorSize(self, arg0: int, arg1: int) -> object: ...
def TensorSizeSignature(self, arg0: int, arg1: int) -> object: ...
def TensorSparsityParameters(self, arg0: int, arg1: int) -> object: ...
def TensorType(self, arg0: int, arg1: int) -> object: ...
def interpreter(self) -> int: ...
def tensor(self, base_object: object, tensor_index: int, subgraph_index: int = ...) -> object: ...
def CreateWrapperFromBuffer(*args, **kwargs): ...
def CreateWrapperFromFile(*args, **kwargs): ...
tensorflow/lite/python/interpreter_wrapper/interpreter_wrapper.cc
+1218
View File
File diff suppressed because it is too large Load Diff
tensorflow/lite/python/interpreter_wrapper/interpreter_wrapper.h
+169
View File
@@ -0,0 +1,169 @@
/* Copyright 2018 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_PYTHON_INTERPRETER_WRAPPER_INTERPRETER_WRAPPER_H_
#define TENSORFLOW_LITE_PYTHON_INTERPRETER_WRAPPER_INTERPRETER_WRAPPER_H_
#include <functional>
#include <memory>
#include <string>
#include <vector>
// Place `<locale>` before <Python.h> to avoid build failures in macOS.
#include <locale>
// The empty line above is on purpose as otherwise clang-format will
// automatically move <Python.h> before <locale>.
#include <Python.h>
#include "tensorflow/lite/core/interpreter.h"
struct TfLiteDelegate;
// We forward declare TFLite classes here to avoid exposing them to SWIG.
namespace tflite {
class MutableOpResolver;
namespace impl {
class FlatBufferModel;
}
namespace interpreter_wrapper {
class PythonErrorReporter;
class InterpreterWrapper {
public:
using Model = impl::FlatBufferModel;
// SWIG caller takes ownership of pointer.
static InterpreterWrapper* CreateWrapperCPPFromFile(
const char* model_path, int op_resolver_id,
const std::vector<std::string>& registerers, std::string* error_msg,
bool preserve_all_tensors, bool disable_delegate_clustering);
static InterpreterWrapper* CreateWrapperCPPFromFile(
const char* model_path, int op_resolver_id,
const std::vector<std::string>& registerers_by_name,
const std::vector<std::function<void(uintptr_t)>>& registerers_by_func,
std::string* error_msg, bool preserve_all_tensors,
bool disable_delegate_clustering, int num_threads,
bool default_delegate_latest_features,
bool compress_quantization_zero_points,
bool disable_delegate_node_fusion = false,
bool force_delegate_node_profiling = false);
// SWIG caller takes ownership of pointer.
static InterpreterWrapper* CreateWrapperCPPFromBuffer(
PyObject* data, int op_resolver_id,
const std::vector<std::string>& registerers, std::string* error_msg,
bool preserve_all_tensors, bool disable_delegate_clustering);
static InterpreterWrapper* CreateWrapperCPPFromBuffer(
PyObject* data, int op_resolver_id,
const std::vector<std::string>& registerers_by_name,
const std::vector<std::function<void(uintptr_t)>>& registerers_by_func,
std::string* error_msg, bool preserve_all_tensors,
bool disable_delegate_clustering, int num_threads,
bool default_delegate_latest_features,
bool compress_quantization_zero_points,
bool disable_delegate_node_fusion = false,
bool force_delegate_node_profiling = false);
~InterpreterWrapper();
PyObject* AllocateTensors(int subgraph_index);
PyObject* Invoke(int subgraph_index);
PyObject* InputIndices() const;
PyObject* OutputIndices() const;
PyObject* ResizeInputTensor(int i, PyObject* value, bool strict,
int subgraph_index);
int NumTensors(int subgraph_index) const;
int NumSubgraphs() const;
std::string TensorName(int tensor_index, int subgraph_index) const;
PyObject* TensorType(int tensor_index, int subgraph_index) const;
PyObject* TensorSize(int tensor_index, int subgraph_index) const;
PyObject* TensorSizeSignature(int tensor_index, int subgraph_index) const;
PyObject* TensorSparsityParameters(int tensor_index,
int subgraph_index) const;
// Deprecated in favor of TensorQuantizationScales, below.
PyObject* TensorQuantization(int tensor_index, int subgraph_index) const;
PyObject* TensorQuantizationParameters(int tensor_index,
int subgraph_index) const;
PyObject* SetTensor(int tensor_index, PyObject* value, int subgraph_index);
PyObject* GetTensor(int tensor_index, int subgraph_index) const;
PyObject* GetSubgraphIndexFromSignature(const char* signature_key);
PyObject* GetSignatureDefs() const;
PyObject* ResetVariableTensors();
int NumNodes() const;
std::string NodeName(int i) const;
PyObject* NodeInputs(int i) const;
PyObject* NodeOutputs(int i) const;
// Returns a reference to tensor index as a numpy array from subgraph. The
// base_object should be the interpreter object providing the memory.
PyObject* tensor(PyObject* base_object, int tensor_index, int subgraph_index);
PyObject* SetNumThreads(int num_threads);
// Adds a delegate to the interpreter.
PyObject* ModifyGraphWithDelegate(TfLiteDelegate* delegate);
// Experimental and subject to change.
//
// Returns a pointer to the underlying interpreter.
Interpreter* interpreter() { return interpreter_.get(); }
private:
// Helper function to construct an `InterpreterWrapper` object.
// It only returns InterpreterWrapper if it can construct an `Interpreter`.
// Otherwise it returns `nullptr`.
static InterpreterWrapper* CreateInterpreterWrapper(
std::unique_ptr<Model> model, int op_resolver_id,
std::unique_ptr<PythonErrorReporter> error_reporter,
const std::vector<std::string>& registerers_by_name,
const std::vector<std::function<void(uintptr_t)>>& registerers_by_func,
std::string* error_msg, bool preserve_all_tensors,
bool disable_delegate_clustering, int num_threads,
bool default_delegate_latest_features,
bool compress_quantization_zero_points,
bool disable_delegate_node_fusion = false,
bool force_delegate_node_profiling = false);
InterpreterWrapper(std::unique_ptr<Model> model,
std::unique_ptr<PythonErrorReporter> error_reporter,
std::unique_ptr<tflite::MutableOpResolver> resolver,
std::unique_ptr<Interpreter> interpreter);
// InterpreterWrapper is not copyable or assignable. We avoid the use of
// InterpreterWrapper() = delete here for SWIG compatibility.
InterpreterWrapper();
InterpreterWrapper(const InterpreterWrapper& rhs);
// Helper function to resize an input tensor.
PyObject* ResizeInputTensorImpl(int i, PyObject* value);
// The public functions which creates `InterpreterWrapper` should ensure all
// these member variables are initialized successfully. Otherwise it should
// report the error and return `nullptr`.
const std::unique_ptr<Model> model_;
const std::unique_ptr<PythonErrorReporter> error_reporter_;
const std::unique_ptr<tflite::MutableOpResolver> resolver_;
const std::unique_ptr<Interpreter> interpreter_;
};
} // namespace interpreter_wrapper
} // namespace tflite
#endif // TENSORFLOW_LITE_PYTHON_INTERPRETER_WRAPPER_INTERPRETER_WRAPPER_H_
tensorflow/lite/python/interpreter_wrapper/interpreter_wrapper_pybind11.cc
+259
View File
@@ -0,0 +1,259 @@
/* Copyright 2020 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.
==============================================================================*/
#include <cstdint>
#include <functional>
#include <stdexcept>
#include <string>
#include <vector>
#include "pybind11/functional.h" // from @pybind11
#include "pybind11/pybind11.h" // from @pybind11
#include "pybind11/pytypes.h" // from @pybind11
#include "pybind11/stl.h" // from @pybind11
#include "tensorflow/lite/python/interpreter_wrapper/interpreter_wrapper.h"
#include "tensorflow/python/lib/core/pybind11_lib.h"
namespace py = pybind11;
using tflite::interpreter_wrapper::InterpreterWrapper;
PYBIND11_MODULE(_pywrap_tensorflow_interpreter_wrapper, m) {
m.doc() = R"pbdoc(
_pywrap_tensorflow_interpreter_wrapper
-----
)pbdoc";
// pybind11 suggests to convert factory functions into constructors, but
// when bytes are provided the wrapper will be confused which
// constructor to call.
m.def("CreateWrapperFromFile",
[](const std::string& model_path, int op_resolver_id,
const std::vector<std::string>& registerers,
bool preserve_all_tensors, bool disable_delegate_clustering) {
std::string error;
auto* wrapper = ::InterpreterWrapper::CreateWrapperCPPFromFile(
model_path.c_str(), op_resolver_id, registerers, &error,
preserve_all_tensors, disable_delegate_clustering);
if (!wrapper) {
throw std::invalid_argument(error);
}
return wrapper;
});
m.def(
"CreateWrapperFromFile",
[](const std::string& model_path, int op_resolver_id,
const std::vector<std::string>& registerers_by_name,
const std::vector<std::function<void(uintptr_t)>>& registerers_by_func,
bool preserve_all_tensors, bool disable_delegate_clustering,
int num_threads, bool default_delegate_latest_features,
bool compress_quantization_zero_points,
bool disable_delegate_node_fusion = false,
bool force_delegate_node_profiling = false) {
std::string error;
auto* wrapper = ::InterpreterWrapper::CreateWrapperCPPFromFile(
model_path.c_str(), op_resolver_id, registerers_by_name,
registerers_by_func, &error, preserve_all_tensors,
disable_delegate_clustering, num_threads,
default_delegate_latest_features, compress_quantization_zero_points,
disable_delegate_node_fusion, force_delegate_node_profiling);
if (!wrapper) {
throw std::invalid_argument(error);
}
return wrapper;
});
m.def("CreateWrapperFromBuffer",
[](const py::bytes& data, int op_resolver_id,
const std::vector<std::string>& registerers,
bool preserve_all_tensors, bool disable_delegate_clustering) {
std::string error;
auto* wrapper = ::InterpreterWrapper::CreateWrapperCPPFromBuffer(
data.ptr(), op_resolver_id, registerers, &error,
preserve_all_tensors, disable_delegate_clustering);
if (!wrapper) {
throw std::invalid_argument(error);
}
return wrapper;
});
m.def(
"CreateWrapperFromBuffer",
[](const py::bytes& data, int op_resolver_id,
const std::vector<std::string>& registerers_by_name,
const std::vector<std::function<void(uintptr_t)>>& registerers_by_func,
bool preserve_all_tensors, bool disable_delegate_clustering,
int num_threads, bool default_delegate_latest_features,
bool compress_quantization_zero_points,
bool disable_delegate_node_fusion = false,
bool force_delegate_node_profiling = false) {
std::string error;
auto* wrapper = ::InterpreterWrapper::CreateWrapperCPPFromBuffer(
data.ptr(), op_resolver_id, registerers_by_name,
registerers_by_func, &error, preserve_all_tensors,
disable_delegate_clustering, num_threads,
default_delegate_latest_features, compress_quantization_zero_points,
disable_delegate_node_fusion, force_delegate_node_profiling);
if (!wrapper) {
throw std::invalid_argument(error);
}
return wrapper;
});
py::class_<InterpreterWrapper>(m, "InterpreterWrapper", py::module_local())
.def(
"AllocateTensors",
[](InterpreterWrapper& self, int subgraph_index) {
return tensorflow::PyoOrThrow(self.AllocateTensors(subgraph_index));
},
// LINT.IfChange
py::arg("subgraph_index") = -1)
// LINT.ThenChange(//tensorflow/lite/python/interpreter_wrapper/interpreter_wrapper.cc)
.def(
"Invoke",
[](InterpreterWrapper& self, int subgraph_index) {
return tensorflow::PyoOrThrow(self.Invoke(subgraph_index));
},
py::arg("subgraph_index") = 0)
.def("InputIndices",
[](const InterpreterWrapper& self) {
return tensorflow::PyoOrThrow(self.InputIndices());
})
.def("OutputIndices",
[](InterpreterWrapper& self) {
return tensorflow::PyoOrThrow(self.OutputIndices());
})
.def(
"ResizeInputTensor",
[](InterpreterWrapper& self, int i, py::handle& value, bool strict,
int subgraph_index) {
return tensorflow::PyoOrThrow(
self.ResizeInputTensor(i, value.ptr(), strict, subgraph_index));
},
py::arg("i"), py::arg("value"), py::arg("strict"),
py::arg("subgraph_index") = 0)
.def("NumTensors", &InterpreterWrapper::NumTensors)
.def("NumSubgraphs", &InterpreterWrapper::NumSubgraphs)
.def("TensorName", &InterpreterWrapper::TensorName)
.def("TensorType",
[](const InterpreterWrapper& self, int tensor_index,
int subgraph_index) {
return tensorflow::PyoOrThrow(
self.TensorType(tensor_index, subgraph_index));
})
.def("TensorSize",
[](const InterpreterWrapper& self, int tensor_index,
int subgraph_index) {
return tensorflow::PyoOrThrow(
self.TensorSize(tensor_index, subgraph_index));
})
.def("TensorSizeSignature",
[](const InterpreterWrapper& self, int tensor_index,
int subgraph_index) {
return tensorflow::PyoOrThrow(
self.TensorSizeSignature(tensor_index, subgraph_index));
})
.def("TensorSparsityParameters",
[](const InterpreterWrapper& self, int tensor_index,
int subgraph_index) {
return tensorflow::PyoOrThrow(
self.TensorSparsityParameters(tensor_index, subgraph_index));
})
.def(
"TensorQuantization",
[](const InterpreterWrapper& self, int tensor_index,
int subgraph_index) {
return tensorflow::PyoOrThrow(
self.TensorQuantization(tensor_index, subgraph_index));
},
R"pbdoc(
Deprecated in favor of TensorQuantizationParameters.
)pbdoc")
.def("TensorQuantizationParameters",
[](InterpreterWrapper& self, int tensor_index, int subgraph_index) {
return tensorflow::PyoOrThrow(self.TensorQuantizationParameters(
tensor_index, subgraph_index));
})
.def(
"SetTensor",
[](InterpreterWrapper& self, int tensor_index, py::handle& value,
int subgraph_index) {
return tensorflow::PyoOrThrow(
self.SetTensor(tensor_index, value.ptr(), subgraph_index));
},
py::arg("i"), py::arg("value"), py::arg("subgraph_index") = 0)
.def(
"GetTensor",
[](const InterpreterWrapper& self, int tensor_index,
int subgraph_index) {
return tensorflow::PyoOrThrow(
self.GetTensor(tensor_index, subgraph_index));
},
py::arg("tensor_index"), py::arg("subgraph_index") = 0)
.def("GetSubgraphIndexFromSignature",
[](InterpreterWrapper& self, const char* signature_key) {
return tensorflow::PyoOrThrow(
self.GetSubgraphIndexFromSignature(signature_key));
})
.def("GetSignatureDefs",
[](InterpreterWrapper& self) {
return tensorflow::PyoOrThrow(self.GetSignatureDefs());
})
.def("ResetVariableTensors",
[](InterpreterWrapper& self) {
return tensorflow::PyoOrThrow(self.ResetVariableTensors());
})
.def("NumNodes", &InterpreterWrapper::NumNodes)
.def("NodeName", &InterpreterWrapper::NodeName)
.def("NodeInputs",
[](const InterpreterWrapper& self, int i) {
return tensorflow::PyoOrThrow(self.NodeInputs(i));
})
.def("NodeOutputs",
[](const InterpreterWrapper& self, int i) {
return tensorflow::PyoOrThrow(self.NodeOutputs(i));
})
.def(
"tensor",
[](InterpreterWrapper& self, py::handle& base_object,
int tensor_index, int subgraph_index) {
return tensorflow::PyoOrThrow(
self.tensor(base_object.ptr(), tensor_index, subgraph_index));
},
R"pbdoc(
Returns a reference to tensor index as a numpy array from subgraph.
The base_object should be the interpreter object providing the
memory.
)pbdoc",
py::arg("base_object"), py::arg("tensor_index"),
py::arg("subgraph_index") = 0)
.def(
"ModifyGraphWithDelegate",
// Address of the delegate is passed as an argument.
[](InterpreterWrapper& self, uintptr_t delegate_ptr) {
return tensorflow::PyoOrThrow(self.ModifyGraphWithDelegate(
reinterpret_cast<TfLiteDelegate*>(delegate_ptr)));
},
R"pbdoc(
Adds a delegate to the interpreter.
)pbdoc")
.def(
"SetNumThreads",
[](InterpreterWrapper& self, int num_threads) {
return tensorflow::PyoOrThrow(self.SetNumThreads(num_threads));
},
R"pbdoc(
ask the interpreter to set the number of threads to use.
)pbdoc")
.def("interpreter", [](InterpreterWrapper& self) {
return reinterpret_cast<intptr_t>(self.interpreter());
});
}
tensorflow/lite/python/interpreter_wrapper/numpy.cc
+327
View File
@@ -0,0 +1,327 @@
/* 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.
==============================================================================*/
#include <cstddef>
#include <cstdint>
#include <memory>
#define TFLITE_IMPORT_NUMPY // See numpy.h for explanation.
#include "tensorflow/lite/core/c/c_api_types.h"
#include "tensorflow/lite/core/c/common.h"
#include "tensorflow/lite/python/interpreter_wrapper/numpy.h"
namespace tflite {
namespace python {
void ImportNumpy() { import_array1(); }
} // namespace python
namespace python_utils {
struct PyObjectDereferencer {
void operator()(PyObject* py_object) const { Py_DECREF(py_object); }
};
using UniquePyObjectRef = std::unique_ptr<PyObject, PyObjectDereferencer>;
int TfLiteTypeToPyArrayType(TfLiteType tf_lite_type) {
switch (tf_lite_type) {
case kTfLiteFloat32:
return NPY_FLOAT32;
case kTfLiteFloat16:
return NPY_FLOAT16;
case kTfLiteBFloat16:
// TODO(b/329491949): Supports other ml_dtypes user-defined types.
return NPY_USERDEF;
case kTfLiteFloat64:
return NPY_FLOAT64;
case kTfLiteInt32:
return NPY_INT32;
case kTfLiteUInt32:
return NPY_UINT32;
case kTfLiteUInt16:
return NPY_UINT16;
case kTfLiteInt16:
return NPY_INT16;
case kTfLiteInt4:
// TODO(b/246806634): NPY_INT4 currently doesn't exist
return NPY_BYTE;
case kTfLiteUInt4:
return NPY_UINT8;
case kTfLiteFloat8E4M3FN:
case kTfLiteFloat8E5M2:
return NPY_BYTE;
case kTfLiteInt2:
// TODO(b/246806634): NPY_INT2 currently doesn't exist
return NPY_BYTE;
case kTfLiteUInt8:
return NPY_UINT8;
case kTfLiteInt8:
return NPY_INT8;
case kTfLiteInt64:
return NPY_INT64;
case kTfLiteUInt64:
return NPY_UINT64;
case kTfLiteString:
return NPY_STRING;
case kTfLiteBool:
return NPY_BOOL;
case kTfLiteComplex64:
return NPY_COMPLEX64;
case kTfLiteComplex128:
return NPY_COMPLEX128;
case kTfLiteResource:
case kTfLiteVariant:
return NPY_OBJECT;
case kTfLiteNoType:
return NPY_NOTYPE;
// Avoid default so compiler errors created when new types are made.
}
return NPY_NOTYPE;
} // NOLINT(direct import ndarraytypes.h cannot be used here)
TfLiteType TfLiteTypeFromPyType(int py_type) {
switch (py_type) {
case NPY_FLOAT32:
return kTfLiteFloat32;
case NPY_FLOAT16:
return kTfLiteFloat16;
case NPY_FLOAT64:
return kTfLiteFloat64;
case NPY_INT32:
return kTfLiteInt32;
case NPY_UINT32:
return kTfLiteUInt32;
case NPY_INT16:
return kTfLiteInt16;
case NPY_UINT16:
return kTfLiteUInt16;
case NPY_UINT8:
return kTfLiteUInt8;
case NPY_INT8:
return kTfLiteInt8;
case NPY_INT64:
return kTfLiteInt64;
case NPY_UINT64:
return kTfLiteUInt64;
case NPY_BOOL:
return kTfLiteBool;
case NPY_OBJECT:
case NPY_STRING:
case NPY_UNICODE:
return kTfLiteString;
case NPY_COMPLEX64:
return kTfLiteComplex64;
case NPY_COMPLEX128:
return kTfLiteComplex128;
case NPY_USERDEF:
// User-defined types are defined in ml_dtypes. (bfloat16, float8, etc.)
// For now, we only support bfloat16.
return kTfLiteBFloat16;
}
return kTfLiteNoType;
} // NOLINT(direct import ndarraytypes.h cannot be used here)
TfLiteType TfLiteTypeFromPyArray(PyArrayObject* array) {
int pyarray_type = PyArray_TYPE(array);
return TfLiteTypeFromPyType(pyarray_type);
}
#if PY_VERSION_HEX >= 0x03030000
bool FillStringBufferFromPyUnicode(PyObject* value,
DynamicBuffer* dynamic_buffer) {
Py_ssize_t len = -1;
const char* buf = PyUnicode_AsUTF8AndSize(value, &len);
if (buf == nullptr) {
PyErr_SetString(PyExc_ValueError, "PyUnicode_AsUTF8AndSize() failed.");
return false;
}
dynamic_buffer->AddString(buf, len);
return true;
}
#else
bool FillStringBufferFromPyUnicode(PyObject* value,
DynamicBuffer* dynamic_buffer) {
UniquePyObjectRef utemp(PyUnicode_AsUTF8String(value));
if (!utemp) {
PyErr_SetString(PyExc_ValueError, "PyUnicode_AsUTF8String() failed.");
return false;
}
char* buf = nullptr;
Py_ssize_t len = -1;
if (PyBytes_AsStringAndSize(utemp.get(), &buf, &len) == -1) {
PyErr_SetString(PyExc_ValueError, "PyBytes_AsStringAndSize() failed.");
return false;
}
dynamic_buffer->AddString(buf, len);
return true;
}
#endif
bool FillStringBufferFromPyString(PyObject* value,
DynamicBuffer* dynamic_buffer) {
if (PyUnicode_Check(value)) {
return FillStringBufferFromPyUnicode(value, dynamic_buffer);
}
char* buf = nullptr;
Py_ssize_t len = -1;
if (PyBytes_AsStringAndSize(value, &buf, &len) == -1) {
PyErr_SetString(PyExc_ValueError, "PyBytes_AsStringAndSize() failed.");
return false;
}
dynamic_buffer->AddString(buf, len);
return true;
}
bool FillStringBufferWithPyArray(PyObject* value,
DynamicBuffer* dynamic_buffer) {
if (!PyArray_Check(value)) {
PyErr_Format(PyExc_ValueError,
"Passed in value type is not a numpy array, got type %s.",
value->ob_type->tp_name);
return false;
}
PyArrayObject* array = reinterpret_cast<PyArrayObject*>(value);
switch (PyArray_TYPE(array)) {
case NPY_OBJECT:
case NPY_STRING:
case NPY_UNICODE: {
if (PyArray_NDIM(array) == 0) {
dynamic_buffer->AddString(static_cast<char*>(PyArray_DATA(array)),
PyArray_NBYTES(array));
return true;
}
UniquePyObjectRef iter(PyArray_IterNew(value));
while (PyArray_ITER_NOTDONE(iter.get())) {
UniquePyObjectRef item(PyArray_GETITEM(
array, reinterpret_cast<char*>(PyArray_ITER_DATA(iter.get()))));
if (!FillStringBufferFromPyString(item.get(), dynamic_buffer)) {
return false;
}
PyArray_ITER_NEXT(iter.get());
}
return true;
}
default:
break;
}
PyErr_Format(PyExc_ValueError,
"Cannot use numpy array of type %d for string tensor.",
PyArray_TYPE(array));
return false;
}
// Helper function to pack int8/uint8 numpy array data into an INT4/UINT4
// tensor.
PyObject* Set4BitTensor(TfLiteTensor* tensor, PyArrayObject* array,
int tensor_index) {
TfLiteType incoming_type = TfLiteTypeFromPyArray(array);
if (tensor->type == kTfLiteInt4) {
if (incoming_type != kTfLiteInt8) {
PyErr_Format(PyExc_ValueError,
"Cannot set tensor:"
" Expected a numpy array of int8 for INT4 input "
"%d, name: %s, but got %s",
tensor_index, tensor->name,
TfLiteTypeGetName(incoming_type));
return nullptr;
}
} else if (tensor->type == kTfLiteUInt4) {
if (incoming_type != kTfLiteUInt8) {
PyErr_Format(PyExc_ValueError,
"Cannot set tensor:"
" Expected a numpy array of uint8 for UINT4 input "
"%d, name: %s, but got %s",
tensor_index, tensor->name,
TfLiteTypeGetName(incoming_type));
return nullptr;
}
}
size_t num_elements = 1;
for (int i = 0; i < tensor->dims->size; ++i) {
num_elements *= tensor->dims->data[i];
}
size_t expected_packed_bytes = (num_elements + 1) / 2;
size_t actual_numpy_bytes = PyArray_NBYTES(array);
const char* tensor_type_name = TfLiteTypeGetName(tensor->type);
if (actual_numpy_bytes != num_elements) {
PyErr_Format(PyExc_ValueError,
"Cannot set tensor:"
" Numpy array for %s input %d, name: %s, has %zu bytes, "
"but expected %zu bytes for %zu elements",
tensor_type_name, tensor_index, tensor->name,
actual_numpy_bytes, num_elements, num_elements);
return nullptr;
}
if (tensor->data.raw == nullptr && tensor->bytes) {
PyErr_Format(PyExc_ValueError,
"Cannot set tensor:"
" Tensor is unallocated. Try calling allocate_tensors()"
" first for input %d, name: %s",
tensor_index, tensor->name);
return nullptr;
}
// Pack the int8/uint8 array into int4/uint4
uint8_t* packed_data = reinterpret_cast<uint8_t*>(tensor->data.raw);
if (tensor->type == kTfLiteInt4) {
int8_t* numpy_data = reinterpret_cast<int8_t*>(PyArray_DATA(array));
for (size_t i = 0; i < expected_packed_bytes; ++i) {
int8_t first_nibble = numpy_data[2 * i];
int8_t second_nibble =
(2 * i + 1 < num_elements) ? numpy_data[2 * i + 1] : 0;
if ((first_nibble < -8 || first_nibble > 7) ||
(second_nibble < -8 || second_nibble > 7)) {
PyErr_Format(PyExc_ValueError,
"Cannot set tensor:"
" Values for INT4 input must be between -8 and 7.");
return nullptr;
}
// Pack the two int8 values into a single byte. The first nibble
// occupies the lower 4 bits and the second nibble occupies the upper 4
// bits. We mask the first nibble with 0x0F to ensure only the lower 4
// bits are used, handling potential sign extension in the int8 value.
packed_data[i] = (first_nibble & 0x0F) | (second_nibble << 4);
}
} else { // kTfLiteUInt4
uint8_t* numpy_data = reinterpret_cast<uint8_t*>(PyArray_DATA(array));
for (size_t i = 0; i < expected_packed_bytes; ++i) {
uint8_t first_nibble = numpy_data[2 * i];
uint8_t second_nibble =
(2 * i + 1 < num_elements) ? numpy_data[2 * i + 1] : 0;
if (first_nibble > 15 || second_nibble > 15) {
PyErr_Format(PyExc_ValueError,
"Cannot set tensor:"
" Values for UINT4 input must be between 0 and 15.");
return nullptr;
}
packed_data[i] = (first_nibble & 0x0F) | (second_nibble << 4);
}
}
Py_RETURN_NONE;
}
} // namespace python_utils
} // namespace tflite
tensorflow/lite/python/interpreter_wrapper/numpy.h
+83
View File
@@ -0,0 +1,83 @@
/* 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_PYTHON_INTERPRETER_WRAPPER_NUMPY_H_
#define TENSORFLOW_LITE_PYTHON_INTERPRETER_WRAPPER_NUMPY_H_
#ifdef PyArray_Type
#error "Numpy cannot be included before numpy.h."
#endif
// Disallow Numpy 1.7 deprecated symbols.
#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
// To handle PyArray_* calles, numpy defines a static lookup table called
// PyArray_API, or PY_ARRAY_UNIQUE_SYMBOL, if defined. This causes the
// PyArray_* pointers to be different for different translation units, unless
// we take care of selectivel defined NO_IMPORT_ARRAY.
//
// Virtually every usage will define NO_IMPORT_ARRAY, and will have access to
// the lookup table via:
// extern void **PyArray_API;
// In numpy.cc we will define TFLITE_IMPORT_NUMPY, effectively disabling that
// and instead using:
// void **PyArray_API;
// which is initialized when ImportNumpy() is called.
//
// If we don't define PY_ARRAY_UNIQUE_SYMBOL then PyArray_API is a static
// variable, which causes strange crashes when the pointers are used across
// translation unit boundaries.
//
// For more info see https://sourceforge.net/p/numpy/mailman/message/5700519
// See also tensorflow/compiler/xla/tsl/python/lib/core/numpy.h for a similar
// approach.
#define PY_ARRAY_UNIQUE_SYMBOL _tflite_numpy_api
#ifndef TFLITE_IMPORT_NUMPY
#define NO_IMPORT_ARRAY
#endif
#include <Python.h>
#include "numpy/arrayobject.h" // IWYU pragma: export
#include "numpy/ufuncobject.h" // IWYU pragma: export
#include "tensorflow/lite/core/c/common.h"
#include "tensorflow/lite/string_util.h"
namespace tflite {
namespace python {
void ImportNumpy();
} // namespace python
namespace python_utils {
int TfLiteTypeToPyArrayType(TfLiteType tf_lite_type);
TfLiteType TfLiteTypeFromPyType(int py_type);
TfLiteType TfLiteTypeFromPyArray(PyArrayObject* array);
bool FillStringBufferWithPyArray(PyObject* value,
DynamicBuffer* dynamic_buffer);
// Helper function to pack int8/uint8 numpy array data into an INT4/UINT4
// tensor.
PyObject* Set4BitTensor(TfLiteTensor* tensor, PyArrayObject* array,
int tensor_index);
} // namespace python_utils
} // namespace tflite
#endif // TENSORFLOW_LITE_PYTHON_INTERPRETER_WRAPPER_NUMPY_H_
tensorflow/lite/python/interpreter_wrapper/python_error_reporter.cc
+49
View File
@@ -0,0 +1,49 @@
/* 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.
==============================================================================*/
#include "tensorflow/lite/python/interpreter_wrapper/python_error_reporter.h"
#include <Python.h>
#include <cstdarg>
#include <cstdio>
#include <string>
namespace tflite {
namespace interpreter_wrapper {
// Report an error message
int PythonErrorReporter::Report(const char* format, va_list args) {
char buf[1024];
int formatted = vsnprintf(buf, sizeof(buf), format, args);
buffer_ << buf;
return formatted;
}
// Set's a Python runtime exception with the last error.
PyObject* PythonErrorReporter::exception() {
std::string last_message = message();
PyErr_SetString(PyExc_RuntimeError, last_message.c_str());
return nullptr;
}
// Gets the last error message and clears the buffer.
std::string PythonErrorReporter::message() {
std::string value = buffer_.str();
buffer_.clear();
return value;
}
} // namespace interpreter_wrapper
} // namespace tflite
tensorflow/lite/python/interpreter_wrapper/python_error_reporter.h
+49
View File
@@ -0,0 +1,49 @@
/* 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_PYTHON_INTERPRETER_WRAPPER_PYTHON_ERROR_REPORTER_H_
#define TENSORFLOW_LITE_PYTHON_INTERPRETER_WRAPPER_PYTHON_ERROR_REPORTER_H_
#include <Python.h>
#include <sstream>
#include <string>
#include "tensorflow/lite/stateful_error_reporter.h"
namespace tflite {
namespace interpreter_wrapper {
class PythonErrorReporter : public tflite::StatefulErrorReporter {
public:
PythonErrorReporter() {}
// Report an error message
int Report(const char* format, va_list args) override;
// Sets a Python runtime exception with the last error and
// clears the error message buffer.
PyObject* exception();
// Gets the last error message and clears the buffer.
std::string message() override;
private:
std::stringstream buffer_;
};
} // namespace interpreter_wrapper
} // namespace tflite
#endif // TENSORFLOW_LITE_PYTHON_INTERPRETER_WRAPPER_PYTHON_ERROR_REPORTER_H_
tensorflow/lite/python/interpreter_wrapper/python_utils.cc
+45
View File
@@ -0,0 +1,45 @@
/* 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.
==============================================================================*/
#include "tensorflow/lite/python/interpreter_wrapper/python_utils.h"
#include <cstddef>
namespace tflite {
namespace python_utils {
int ConvertFromPyString(PyObject* obj, char** data, Py_ssize_t* length) {
#if PY_MAJOR_VERSION >= 3
if (PyUnicode_Check(obj)) {
// const_cast<> is for CPython 3.7 finally adding const to the API.
*data = const_cast<char*>(PyUnicode_AsUTF8AndSize(obj, length));
return *data == nullptr ? -1 : 0;
}
return PyBytes_AsStringAndSize(obj, data, length);
#else
return PyString_AsStringAndSize(obj, data, length);
#endif
}
PyObject* ConvertToPyString(const char* data, size_t length) {
#if PY_MAJOR_VERSION >= 3
return PyBytes_FromStringAndSize(data, length);
#else
return PyString_FromStringAndSize(data, length);
#endif
}
} // namespace python_utils
} // namespace tflite
tensorflow/lite/python/interpreter_wrapper/python_utils.h
+35
View File
@@ -0,0 +1,35 @@
/* 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_PYTHON_INTERPRETER_WRAPPER_PYTHON_UTILS_H_
#define TENSORFLOW_LITE_PYTHON_INTERPRETER_WRAPPER_PYTHON_UTILS_H_
#include <Python.h>
#include <cstddef>
namespace tflite {
namespace python_utils {
struct PyDecrefDeleter {
void operator()(PyObject* p) const { Py_DECREF(p); }
};
int ConvertFromPyString(PyObject* obj, char** data, Py_ssize_t* length);
PyObject* ConvertToPyString(const char* data, size_t length);
} // namespace python_utils
} // namespace tflite
#endif // TENSORFLOW_LITE_PYTHON_INTERPRETER_WRAPPER_PYTHON_UTILS_H_
tensorflow/lite/python/kernel_tests/signal/BUILD
+38
View File
@@ -0,0 +1,38 @@
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
licenses = ["notice"],
)
py_library(
name = "test_util",
srcs = ["test_util.py"],
strict_deps = True,
deps = [
"//tensorflow/lite/python:interpreter",
"//tensorflow/lite/python:lite",
"//tensorflow/python/eager:def_function",
],
)
cuda_py_strict_test(
name = "window_ops_test",
srcs = ["window_ops_test.py"],
shard_count = 4,
tags = [
"no_rocm",
"no_windows_gpu",
],
deps = [
":test_util",
"//tensorflow/python/framework:for_generated_wrappers",
"//tensorflow/python/framework:tensor_spec",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops/signal:window_ops",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
)
tensorflow/lite/python/kernel_tests/signal/test_util.py
+69
View File
@@ -0,0 +1,69 @@
# 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.
# ==============================================================================
"""Test utilities for tf.signal."""
from tensorflow.lite.python import interpreter
from tensorflow.lite.python import lite
from tensorflow.python.eager import def_function
def tflite_convert(fn, input_templates):
"""Converts the provided fn to tf.lite model.
Args:
fn: A callable that expects a list of inputs like input_templates that
returns a tensor or structure of tensors.
input_templates: A list of Tensors, ndarrays or TensorSpecs describing the
inputs that fn expects. The actual values of the Tensors or ndarrays are
unused.
Returns:
The serialized tf.lite model.
"""
fn = def_function.function(fn)
concrete_func = fn.get_concrete_function(*input_templates)
converter = lite.TFLiteConverterV2([concrete_func])
return converter.convert()
def evaluate_tflite_model(tflite_model, input_ndarrays):
"""Evaluates the provided tf.lite model with the given input ndarrays.
Args:
tflite_model: bytes. The serialized tf.lite model.
input_ndarrays: A list of NumPy arrays to feed as input to the model.
Returns:
A list of ndarrays produced by the model.
Raises:
ValueError: If the number of input arrays does not match the number of
inputs the model expects.
"""
the_interpreter = interpreter.Interpreter(model_content=tflite_model)
the_interpreter.allocate_tensors()
input_details = the_interpreter.get_input_details()
output_details = the_interpreter.get_output_details()
if len(input_details) != len(input_ndarrays):
raise ValueError('Wrong number of inputs: provided=%s, '
'input_details=%s output_details=%s' % (
input_ndarrays, input_details, output_details))
for input_tensor, data in zip(input_details, input_ndarrays):
the_interpreter.set_tensor(input_tensor['index'], data)
the_interpreter.invoke()
return [the_interpreter.get_tensor(details['index'])
for details in output_details]
tensorflow/lite/python/kernel_tests/signal/window_ops_test.py
+60
View File
@@ -0,0 +1,60 @@
# 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.
# ==============================================================================
"""Tests for window_ops."""
from absl.testing import parameterized
import numpy as np
from tensorflow.lite.python.kernel_tests.signal import test_util
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import tensor_spec
from tensorflow.python.framework import test_util as tf_test_util
from tensorflow.python.ops.signal import window_ops
from tensorflow.python.platform import test
@tf_test_util.run_all_in_graph_and_eager_modes
class WindowOpsTest(test.TestCase, parameterized.TestCase):
@parameterized.parameters(
# Only float32 is supported.
(window_ops.hann_window, 10, False, dtypes.float32),
(window_ops.hann_window, 10, True, dtypes.float32),
(window_ops.hamming_window, 10, False, dtypes.float32),
(window_ops.hamming_window, 10, True, dtypes.float32),
(window_ops.vorbis_window, 12, None, dtypes.float32),
)
def test_tflite_convert(self, window_fn, window_length, periodic, dtype):
def fn(window_length):
try:
return window_fn(window_length, periodic=periodic, dtype=dtype)
except TypeError:
return window_fn(window_length, dtype=dtype)
tflite_model = test_util.tflite_convert(
fn, [tensor_spec.TensorSpec(shape=[], dtype=dtypes.int32)]
)
window_length = np.array(window_length).astype(np.int32)
(actual_output,) = test_util.evaluate_tflite_model(
tflite_model, [window_length]
)
expected_output = self.evaluate(fn(window_length))
self.assertAllClose(actual_output, expected_output, rtol=1e-6, atol=1e-6)
if __name__ == '__main__':
test.main()
tensorflow/lite/python/lite.py
+3455
View File
File diff suppressed because it is too large Load Diff
tensorflow/lite/python/lite_constants.py
+83
View File
@@ -0,0 +1,83 @@
# Copyright 2018 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.
# ==============================================================================
"""Constants for TFLite."""
from tensorflow.compiler.mlir.lite import converter_flags_pb2 as _converter_flags_pb2
from tensorflow.python.framework import dtypes
from tensorflow.python.util.all_util import remove_undocumented
from tensorflow.python.util.tf_export import tf_export as _tf_export
FLOAT = dtypes.float32
FLOAT16 = dtypes.float16
INT32 = dtypes.int32
INT64 = dtypes.int64
STRING = dtypes.string
QUANTIZED_UINT8 = dtypes.uint8
INT8 = dtypes.int8
INT16 = dtypes.int16
COMPLEX64 = dtypes.complex64
TENSORFLOW_GRAPHDEF = _converter_flags_pb2.TENSORFLOW_GRAPHDEF
TFLITE = _converter_flags_pb2.TFLITE
GRAPHVIZ_DOT = _converter_flags_pb2.GRAPHVIZ_DOT
UNSET = _converter_flags_pb2.ConverterFlags.ModelOriginFramework.Name(
_converter_flags_pb2.ConverterFlags.UNSET
)
TENSORFLOW = _converter_flags_pb2.ConverterFlags.ModelOriginFramework.Name(
_converter_flags_pb2.ConverterFlags.TENSORFLOW
)
KERAS = _converter_flags_pb2.ConverterFlags.ModelOriginFramework.Name(
_converter_flags_pb2.ConverterFlags.KERAS
)
JAX = _converter_flags_pb2.ConverterFlags.ModelOriginFramework.Name(
_converter_flags_pb2.ConverterFlags.JAX
)
PYTORCH = _converter_flags_pb2.ConverterFlags.ModelOriginFramework.Name(
_converter_flags_pb2.ConverterFlags.PYTORCH
)
_tf_export(v1=["lite.constants.FLOAT"]).export_constant(__name__, "FLOAT")
_tf_export(v1=["lite.constants.FLOAT16"]).export_constant(__name__, "FLOAT16")
_tf_export(v1=["lite.constants.INT32"]).export_constant(__name__, "INT32")
_tf_export(v1=["lite.constants.INT64"]).export_constant(__name__, "INT64")
_tf_export(v1=["lite.constants.STRING"]).export_constant(__name__, "STRING")
_tf_export(v1=["lite.constants.QUANTIZED_UINT8"]).export_constant(
__name__, "QUANTIZED_UINT8")
_tf_export(v1=["lite.constants.INT8"]).export_constant(__name__, "INT8")
_tf_export(v1=["lite.constants.INT16"]).export_constant(__name__, "INT16")
_tf_export(v1=["lite.constants.TFLITE"]).export_constant(__name__, "TFLITE")
_tf_export(v1=["lite.constants.GRAPHVIZ_DOT"]).export_constant(
__name__, "GRAPHVIZ_DOT")
_allowed_symbols = [
"FLOAT",
"FLOAT16",
"INT32",
"INT64",
"STRING",
"QUANTIZED_UINT8",
"INT8",
"INT16",
"COMPLEX64",
"TENSORFLOW_GRAPHDEF",
"TFLITE",
"GRAPHVIZ_DOT",
"UNSET",
"TENSORFLOW",
"KERAS",
"JAX",
"PYTORCH",
]
remove_undocumented(__name__, _allowed_symbols)
tensorflow/lite/python/lite_flex_test.py
+445
View File
@@ -0,0 +1,445 @@
# 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.
# ==============================================================================
"""Tests for lite.py functionality related to select TF op usage."""
import os
from absl.testing import parameterized
import numpy as np
import tensorflow as tf
from tensorflow.core.framework import graph_pb2
from tensorflow.lite.python import lite
from tensorflow.lite.python import test_util as tflite_test_util
from tensorflow.lite.python.convert import register_custom_opdefs
from tensorflow.lite.python.interpreter import Interpreter
from tensorflow.lite.python.testdata import double_op
from tensorflow.python.client import session
from tensorflow.python.eager import def_function
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import test_util
from tensorflow.python.framework.importer import import_graph_def
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import list_ops
from tensorflow.python.ops import nn_ops
from tensorflow.python.ops import variables
from tensorflow.python.platform import test
from tensorflow.python.saved_model import saved_model
from tensorflow.python.trackable import autotrackable
class FromSessionTest(test_util.TensorFlowTestCase, parameterized.TestCase):
@parameterized.named_parameters(
('EnableMlirConverter', True), # enable mlir
('DisableMlirConverter', False)) # disable mlir
def testFlexMode(self, enable_mlir):
with ops.Graph().as_default():
in_tensor = array_ops.placeholder(shape=[1, 4], dtype=dtypes.float32)
out_tensor = in_tensor + in_tensor
sess = session.Session()
# Convert model and ensure model is not None.
converter = lite.TFLiteConverter.from_session(sess, [in_tensor],
[out_tensor])
converter.target_spec.supported_ops = set([lite.OpsSet.SELECT_TF_OPS])
converter.experimental_new_converter = enable_mlir
tflite_model = converter.convert()
self.assertTrue(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.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 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 testDeprecatedFlags(self):
with ops.Graph().as_default():
in_tensor = array_ops.placeholder(shape=[1, 4], dtype=dtypes.float32)
out_tensor = in_tensor + in_tensor
sess = session.Session()
# Convert model and ensure model is not None.
converter = lite.TFLiteConverter.from_session(sess, [in_tensor],
[out_tensor])
converter.target_ops = set([lite.OpsSet.SELECT_TF_OPS])
# Ensure `target_ops` is set to the correct value after flag deprecation.
self.assertEqual(converter.target_ops, set([lite.OpsSet.SELECT_TF_OPS]))
self.assertEqual(converter.target_spec.supported_ops,
set([lite.OpsSet.SELECT_TF_OPS]))
tflite_model = converter.convert()
self.assertTrue(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.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())
class FromConcreteFunctionTest(test_util.TensorFlowTestCase,
parameterized.TestCase):
@parameterized.named_parameters(
('EnableMlirConverter', True), # enable mlir
('DisableMlirConverter', False)) # disable mlir
@test_util.run_v2_only
def testFloat(self, enable_mlir):
input_data = constant_op.constant(1., shape=[1])
root = autotrackable.AutoTrackable()
root.v1 = variables.Variable(3.)
root.v2 = variables.Variable(2.)
root.f = def_function.function(lambda x: root.v1 * root.v2 * x)
concrete_func = root.f.get_concrete_function(input_data)
# Convert model.
converter = lite.TFLiteConverterV2.from_concrete_functions([concrete_func],
root)
converter.target_spec.supported_ops = set([lite.OpsSet.SELECT_TF_OPS])
converter.experimental_new_converter = enable_mlir
tflite_model = converter.convert()
# 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([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([24.0], dtype=np.float32)
output_data = interpreter.get_tensor(output_details[0]['index'])
self.assertTrue((expected_output == output_data).all())
# Ensure that input TFLite buffer is not reused for ops such as
# `TensorListSetItem`. The example model has a while loop, and the while body
# has a `TensorListSetItem` op which takes the output from a `Where` op.
@test_util.run_v2_only
def testDisableFlexTensorMemoryReusing(self):
@tf.function(input_signature=[
tf.TensorSpec(shape=[2, 3], dtype=tf.float32, name='x')
])
def model(x):
l = list_ops.tensor_list_reserve(
element_dtype=tf.int64, element_shape=[None, 1], num_elements=2)
init_state = (0, x, l)
condition = lambda i, x, l: i < 2
def body(i, x, l):
element = tf.where(x[i])
l = list_ops.tensor_list_set_item(l, i, element)
return i + 1, x, l
_, _, l_final = tf.while_loop(condition, body, init_state)
return list_ops.tensor_list_stack(l_final, element_dtype=tf.int64)
# Convert model.
converter = lite.TFLiteConverterV2.from_concrete_functions(
[model.get_concrete_function()])
converter.target_spec.supported_ops = set(
[lite.OpsSet.TFLITE_BUILTINS, lite.OpsSet.SELECT_TF_OPS])
tflite_model = converter.convert()
# Check the model produces correct result.
interpreter = Interpreter(model_content=tflite_model)
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
test_input = np.array([[1.0, 2.0, 0.0], [0.0, 5.0, 6.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([0, 1, 1, 2], dtype=np.int64)
output_data = interpreter.get_tensor(output_details[0]['index'])
self.assertTrue((expected_output == np.ndarray.flatten(output_data)).all())
class WithCustomOpTest(test_util.TensorFlowTestCase, parameterized.TestCase):
def _createGraphWithCustomOp(self, opname='CustomAdd'):
custom_opdefs_str = (
'name: \'' + opname + '\' input_arg: {name: \'Input1\' type: DT_FLOAT} '
'input_arg: {name: \'Input2\' type: DT_FLOAT} output_arg: {name: '
'\'Output\' type: DT_FLOAT}')
# Create a graph that has one add op.
new_graph = graph_pb2.GraphDef()
with ops.Graph().as_default():
with session.Session() as sess:
in_tensor = array_ops.placeholder(
shape=[1, 16, 16, 3], dtype=dtypes.float32, name='input')
out_tensor = in_tensor + in_tensor
inputs = {'x': in_tensor}
outputs = {'z': out_tensor}
new_graph.CopyFrom(sess.graph_def)
# Rename Add op name to opname.
for node in new_graph.node:
if node.op.startswith('Add'):
node.op = opname
del node.attr['T']
# Register custom op defs to import modified graph def.
register_custom_opdefs([custom_opdefs_str])
return (new_graph, inputs, outputs)
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 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())
class FromSavedModelTest(test_util.TensorFlowTestCase):
@test_util.run_v2_only
def testFlexResourceVariables(self):
class Model(tf.Module):
def __init__(self):
self.v = tf.Variable([[0.0, 0.0, 0.0, 0.0]])
@tf.function(
input_signature=[tf.TensorSpec(shape=[1, 4], dtype=tf.float32)])
def eval(self, x):
# Control flow is needed to generate "FlexReadVariableOp".
if tf.reduce_mean(x) > 1.0:
self.v.assign_add([[1.0, 1.0, 1.0, 1.0]])
return self.v + x
m = Model()
to_save = m.eval.get_concrete_function()
save_dir = os.path.join(self.get_temp_dir(), 'saved_model')
tf.saved_model.save(m, save_dir, to_save)
converter = lite.TFLiteConverterV2.from_saved_model(save_dir)
converter.target_spec.supported_ops = [
lite.OpsSet.TFLITE_BUILTINS,
lite.OpsSet.SELECT_TF_OPS,
]
converter.experimental_enable_resource_variables = True
tflite_model = converter.convert()
# Check the model works with TensorFlow ops.
interpreter = Interpreter(model_content=tflite_model)
signature_runner = interpreter.get_signature_runner()
outputs = signature_runner(
x=np.array([[1.0, 2.0, 3.0, 4.0]], dtype=np.float32))
expected_output = np.array([[2.0, 3.0, 4.0, 5.0]], dtype=np.float32)
self.assertTrue((expected_output == list(outputs.values())[0]).all)
# Second run.
outputs = signature_runner(
x=np.array([[1.0, 2.0, 3.0, 4.0]], dtype=np.float32))
expected_output = np.array([[3.0, 4.0, 5.0, 6.0]], dtype=np.float32)
self.assertTrue((expected_output == list(outputs.values())[0]).all)
class TFQuantizationTest(test_util.TensorFlowTestCase, parameterized.TestCase):
@parameterized.named_parameters(('DefaultMode', 'DEFAULT'),
('LegacyIntegerMode', 'LEGACY_INTEGER'))
def testAddOp(self, tf_quantization_mode):
root = autotrackable.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())
@parameterized.named_parameters(('DefaultMode', 'DEFAULT'),
('LegacyIntegerMode', 'LEGACY_INTEGER'))
def testL2LossOp(self, tf_quantization_mode):
root = autotrackable.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())
@parameterized.named_parameters(('DefaultMode', 'DEFAULT'),
('LegacyIntegerMode', 'LEGACY_INTEGER'))
def testConvOpWithBias(self, tf_quantization_mode):
class ConvModel(autotrackable.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())
if __name__ == '__main__':
test.main()
tensorflow/lite/python/lite_test.py
+2908
View File
File diff suppressed because it is too large Load Diff
tensorflow/lite/python/lite_v2_test.py
+5585
View File
File diff suppressed because it is too large Load Diff
tensorflow/lite/python/lite_v2_test_util.py
+295
View File
@@ -0,0 +1,295 @@
# 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.
# ==============================================================================
"""Tests for lite.py functionality related to TensorFlow 2.0."""
import os
from absl.testing import parameterized
import numpy as np
import tensorflow as tf
from tensorflow.lite.python.interpreter import Interpreter
from tensorflow.python.eager import def_function
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import tensor_spec
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import variables
from tensorflow.python.trackable import autotrackable
class ModelTest(test_util.TensorFlowTestCase, parameterized.TestCase):
"""Base test class for TensorFlow Lite 2.x model tests."""
def _evaluateTFLiteModel(self, tflite_model, input_data, input_shapes=None):
"""Evaluates the model on the `input_data`.
Args:
tflite_model: TensorFlow Lite model.
input_data: List of EagerTensor const ops containing the input data for
each input tensor.
input_shapes: List of tuples representing the `shape_signature` and the
new shape of each input tensor that has unknown dimensions.
Returns:
[np.ndarray]
"""
interpreter = Interpreter(model_content=tflite_model)
input_details = interpreter.get_input_details()
if input_shapes:
for idx, (shape_signature, final_shape) in enumerate(input_shapes):
self.assertTrue(
(input_details[idx]['shape_signature'] == shape_signature).all())
index = input_details[idx]['index']
interpreter.resize_tensor_input(index, final_shape, strict=True)
interpreter.allocate_tensors()
output_details = interpreter.get_output_details()
input_details = interpreter.get_input_details()
for input_tensor, tensor_data in zip(input_details, input_data):
interpreter.set_tensor(input_tensor['index'], tensor_data.numpy())
interpreter.invoke()
return [
interpreter.get_tensor(details['index']) for details in output_details
]
def _evaluateTFLiteModelUsingSignatureDef(self, tflite_model, signature_key,
inputs):
"""Evaluates the model on the `inputs`.
Args:
tflite_model: TensorFlow Lite model.
signature_key: Signature key.
inputs: Map from input tensor names in the SignatureDef to tensor value.
Returns:
Dictionary of outputs.
Key is the output name in the SignatureDef 'signature_key'
Value is the output value
"""
interpreter = Interpreter(model_content=tflite_model)
signature_runner = interpreter.get_signature_runner(signature_key)
return signature_runner(**inputs)
def _getSimpleVariableModel(self):
root = autotrackable.AutoTrackable()
root.v1 = variables.Variable(3.)
root.v2 = variables.Variable(2.)
root.f = def_function.function(lambda x: root.v1 * root.v2 * x)
return root
def _getSimpleModelWithVariables(self):
class SimpleModelWithOneVariable(autotrackable.AutoTrackable):
"""Basic model with 1 variable."""
def __init__(self):
super(SimpleModelWithOneVariable, self).__init__()
self.var = variables.Variable(array_ops.zeros((1, 10), name='var'))
@def_function.function
def assign_add(self, x):
self.var.assign_add(x)
return self.var
return SimpleModelWithOneVariable()
def _getMultiFunctionModel(self):
class BasicModel(autotrackable.AutoTrackable):
"""Basic model with multiple functions."""
def __init__(self):
self.y = None
self.z = None
@def_function.function
def add(self, x):
if self.y is None:
self.y = variables.Variable(2.)
return x + self.y
@def_function.function
def sub(self, x):
if self.z is None:
self.z = variables.Variable(3.)
return x - self.z
@def_function.function
def mul_add(self, x, y):
if self.z is None:
self.z = variables.Variable(3.)
return x * self.z + y
return BasicModel()
def _getMultiFunctionModelWithSharedWeight(self):
class BasicModelWithSharedWeight(autotrackable.AutoTrackable):
"""Model with multiple functions and a shared weight."""
def __init__(self):
self.weight = constant_op.constant([1.0],
shape=(1, 512, 512, 1),
dtype=dtypes.float32)
@def_function.function
def add(self, x):
return x + self.weight
@def_function.function
def sub(self, x):
return x - self.weight
@def_function.function
def mul(self, x):
return x * self.weight
return BasicModelWithSharedWeight()
def _getMatMulModelWithSmallWeights(self):
class MatMulModelWithSmallWeights(autotrackable.AutoTrackable):
"""MatMul model with small weights and relatively large biases."""
def __init__(self):
self.weight = constant_op.constant([[1e-3, -1e-3], [-2e-4, 2e-4]],
shape=(2, 2),
dtype=dtypes.float32)
self.bias = constant_op.constant([1.28, 2.55],
shape=(2,),
dtype=dtypes.float32)
@def_function.function
def matmul(self, x):
return x @ self.weight + self.bias
return MatMulModelWithSmallWeights()
def _getCeilModel(self):
"""Returns a model with only one ceil op, to test non-quantizable op."""
@def_function.function(input_signature=[
tensor_spec.TensorSpec(shape=(1, 10), dtype=dtypes.float32)
])
def ceil(x):
return math_ops.ceil(x)
def calibration_gen():
for _ in range(5):
yield [np.random.uniform(0, 16, size=(1, 10)).astype(np.float32)]
return ceil, calibration_gen
def _assertValidDebugInfo(self, debug_info):
"""Verify the DebugInfo is valid."""
file_names = set()
for file_path in debug_info.files:
file_names.add(os.path.basename(file_path))
# To make the test independent on how the nodes are created, we only assert
# the name of this test file.
self.assertIn('lite_v2_test.py', file_names)
self.assertNotIn('lite_test.py', file_names)
def _createV2QATLowBitKerasModel(self, shape, weight_only, num_bits, bit_min,
bit_max):
"""Creates a simple QAT num_bits-Weight Keras Model."""
input_name = 'input'
output_name = 'scores'
class ConvWrapper(tf.keras.layers.Wrapper):
"""A Wrapper for simulating QAT on Conv2D layers."""
def build(self, input_shape):
if not self.layer.built:
self.layer.build(input_shape)
self.quantized_weights = self.layer.kernel
def call(self, inputs):
self.layer.kernel = (
tf.quantization.fake_quant_with_min_max_vars_per_channel(
self.quantized_weights, min=[bit_min], max=[bit_max],
num_bits=num_bits, narrow_range=True))
if not weight_only:
quant_inputs = tf.quantization.fake_quant_with_min_max_vars(
inputs, min=0, max=6, num_bits=8)
outputs = self.layer.call(quant_inputs)
return tf.quantization.fake_quant_with_min_max_vars(
outputs, min=0, max=6, num_bits=8)
return self.layer.call(inputs)
input_tensor = tf.keras.layers.Input(shape, name=input_name)
kernel_shape = (shape[-1], 3, 3, 1)
# Ensure constant weights contains the min and max.
initial_weights = np.linspace(
bit_min, bit_max, np.prod(kernel_shape)).reshape(kernel_shape)
test_initializer = tf.constant_initializer(initial_weights)
x = ConvWrapper(tf.keras.layers.Conv2D(
1, (3, 3), kernel_initializer=test_initializer,
activation='relu6'))(input_tensor)
scores = tf.keras.layers.Flatten(name=output_name)(x)
model = tf.keras.Model(input_tensor, scores)
return model, input_name, output_name
def _createReadAssignModel(self, number_of_states=2):
dtype = float
class ReadAssign(tf.keras.layers.Layer):
"""ReadAssign model for the variable quantization test."""
def __init__(self, number_of_states=2, **kwargs):
super().__init__(**kwargs)
self.number_of_states = number_of_states
def build(self, input_shape):
super().build(input_shape)
state_shape = (1, 2, 3)
self.states = [None] * self.number_of_states
for i in range(self.number_of_states):
self.states[i] = self.add_weight(
name=f'states{i}',
shape=state_shape,
trainable=False,
initializer=tf.zeros_initializer,
dtype=dtype,
)
def call(self, inputs):
for state in self.states:
memory = tf.keras.backend.concatenate([state, inputs], 1)
new_state = memory[:, : state.shape[1], :]
state.assign(new_state)
return inputs
def calibration_gen():
for _ in range(5):
yield [np.random.uniform(-1, 1, size=(1, 2, 3)).astype(np.float32)]
inputs = tf.keras.layers.Input(shape=(2, 3), batch_size=1, dtype=dtype)
outputs = ReadAssign(number_of_states)(inputs)
model = tf.keras.Model(inputs, outputs)
return model, calibration_gen
def _getInfFloatModel(self):
root = autotrackable.AutoTrackable()
root.v = constant_op.constant([np.inf], shape=(), dtype=dtypes.float32)
root.f = def_function.function(lambda x: root.v)
return root
tensorflow/lite/python/metrics/BUILD
+158
View File
@@ -0,0 +1,158 @@
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load("//tensorflow:pytype.default.bzl", "pytype_strict_library")
load("//tensorflow:tensorflow.bzl", "py_test")
load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable", "if_portable", "pybind_extension")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
default_visibility = ["//visibility:public"],
licenses = ["notice"],
)
cc_library(
name = "metrics_wrapper_lib",
srcs = if_portable(
if_false = ["wrapper/metrics_wrapper_nonportable.cc"],
if_true = ["wrapper/metrics_wrapper_portable.cc"],
),
hdrs = ["wrapper/metrics_wrapper.h"],
compatible_with = get_compatible_with_portable(),
visibility = ["//visibility:private"],
deps = [
"@xla//third_party/python_runtime:headers",
] + if_portable(
if_false = [
"//learning/brain/google/monitoring:metrics_exporter",
],
if_true = [],
),
)
pybind_extension(
name = "_pywrap_tensorflow_lite_metrics_wrapper",
srcs = ["wrapper/metrics_wrapper_pybind11.cc"],
hdrs = ["wrapper/metrics_wrapper.h"],
common_lib_packages = [
"litert/python",
"tensorflow/lite/python",
],
compatible_with = get_compatible_with_portable(),
enable_stub_generation = True,
pytype_srcs = [
"_pywrap_tensorflow_lite_metrics_wrapper.pyi",
],
visibility = ["//visibility:public"],
wrap_py_init = True,
deps = [
":metrics_wrapper_lib",
"//tensorflow/python/lib/core:pybind11_lib",
"@com_google_protobuf//:protobuf",
"@pybind11",
"@xla//third_party/python_runtime:headers",
],
)
pytype_strict_library(
name = "metrics_wrapper",
srcs = ["wrapper/metrics_wrapper.py"],
deps = [
":_pywrap_tensorflow_lite_metrics_wrapper",
"//tensorflow/compiler/mlir/lite/metrics:converter_error_data_proto_py",
"//tensorflow/compiler/mlir/lite/python:wrap_converter",
],
)
py_test(
name = "metrics_wrapper_test",
srcs = ["wrapper/metrics_wrapper_test.py"],
strict_deps = True,
deps = [
":metrics_wrapper",
#internal proto upb dep
"//tensorflow:tensorflow_py",
"//tensorflow/lite/python:convert",
"//tensorflow/lite/python:lite",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/platform:client_testlib",
],
)
pytype_strict_library(
name = "metrics_interface",
srcs = ["metrics_interface.py"],
compatible_with = get_compatible_with_portable(),
visibility = ["//tensorflow/lite/tools/pip_package:__subpackages__"],
)
genrule(
name = "metrics_py_gen",
srcs = if_portable(
if_false = ["metrics_nonportable.py"],
if_true = ["metrics_portable.py"],
),
outs = ["metrics.py"],
cmd = (
"cat $(SRCS) > $(OUTS)"
),
compatible_with = get_compatible_with_portable(),
)
pytype_strict_library(
name = "metrics",
srcs = ["metrics.py"],
compatible_with = get_compatible_with_portable(),
visibility = ["//tensorflow/lite:__subpackages__"],
deps = if_portable(
if_false = [
"//tensorflow/compiler/mlir/lite/metrics:converter_error_data_proto_py",
":metrics_wrapper",
"//tensorflow/python/eager:monitoring",
],
if_true = [],
) + [":metrics_interface"],
)
py_test(
name = "metrics_test",
srcs = if_portable(
if_false = ["metrics_nonportable_test.py"],
if_true = ["metrics_portable_test.py"],
),
data = [
"//tensorflow/lite/python/testdata/control_flow_v1_saved_model:saved_model.pb",
],
main = if_portable(
if_false = "metrics_nonportable_test.py",
if_true = "metrics_portable_test.py",
),
strict_deps = True,
tags = ["notap"], # TODO(b/373657707): Remove once we debug the failure.
deps = [
":metrics",
"@absl_py//absl/testing:parameterized",
#internal proto upb dep
"//third_party/py/numpy",
"//tensorflow:tensorflow_py",
"//tensorflow/compiler/mlir/lite/metrics:converter_error_data_proto_py",
"//tensorflow/core:protos_all_py",
"//tensorflow/lite/python:convert",
"//tensorflow/lite/python:lite",
"//tensorflow/python/client:session",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:monitoring",
"//tensorflow/python/framework:convert_to_constants",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:importer",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:array_ops_stack",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:string_ops",
"//tensorflow/python/ops/ragged:ragged_tensor",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/platform:resource_loader",
"//tensorflow/python/saved_model",
"//tensorflow/python/trackable:autotrackable",
],
)
tensorflow/lite/python/metrics/_pywrap_tensorflow_lite_metrics_wrapper.pyi
+18
View File
@@ -0,0 +1,18 @@
# Copyright 2023 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.
# ==============================================================================
class MetricsWrapper:
def __init__(self, arg0: str) -> None: ...
def ExportMetrics(self) -> object: ...
tensorflow/lite/python/metrics/metrics_interface.py
+48
View File
@@ -0,0 +1,48 @@
# Copyright 2021 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.
# ==============================================================================
"""Python TFLite metrics helper interface."""
import abc
class TFLiteMetricsInterface(metaclass=abc.ABCMeta):
"""Abstract class for TFLiteMetrics."""
@abc.abstractmethod
def increase_counter_debugger_creation(self):
raise NotImplementedError
@abc.abstractmethod
def increase_counter_interpreter_creation(self):
raise NotImplementedError
@abc.abstractmethod
def increase_counter_converter_attempt(self):
raise NotImplementedError
@abc.abstractmethod
def increase_counter_converter_success(self):
raise NotImplementedError
@abc.abstractmethod
def set_converter_param(self, name, value):
raise NotImplementedError
@abc.abstractmethod
def set_converter_error(self, error_data):
raise NotImplementedError
@abc.abstractmethod
def set_converter_latency(self, value):
raise NotImplementedError
tensorflow/lite/python/metrics/metrics_nonportable.py
+130
View File
@@ -0,0 +1,130 @@
# Copyright 2021 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.
# ==============================================================================
"""Python TFLite metrics helper."""
from typing import Optional, Text
import uuid
from tensorflow.compiler.mlir.lite.metrics import converter_error_data_pb2
from tensorflow.lite.python.metrics import metrics_interface
from tensorflow.lite.python.metrics.wrapper import metrics_wrapper
from tensorflow.python.eager import monitoring
_counter_debugger_creation = monitoring.Counter(
'/tensorflow/lite/quantization_debugger/created',
'Counter for the number of debugger created.')
_counter_interpreter_creation = monitoring.Counter(
'/tensorflow/lite/interpreter/created',
'Counter for number of interpreter created in Python.', 'language')
# The following are conversion metrics. Attempt and success are kept separated
# instead of using a single metric with a label because the converter may
# raise exceptions if conversion failed. That may lead to cases when we are
# unable to capture the conversion attempt. Increasing attempt count at the
# beginning of conversion process and the success count at the end is more
# suitable in these cases.
_counter_conversion_attempt = monitoring.Counter(
'/tensorflow/lite/convert/attempt',
'Counter for number of conversion attempts.')
_counter_conversion_success = monitoring.Counter(
'/tensorflow/lite/convert/success',
'Counter for number of successful conversions.')
_gauge_conversion_params = monitoring.StringGauge(
'/tensorflow/lite/convert/params',
'Gauge for keeping conversion parameters.', 'name')
_gauge_conversion_errors = monitoring.StringGauge(
'/tensorflow/lite/convert/errors',
'Gauge for collecting conversion errors. The value represents the error '
'message.', 'component', 'subcomponent', 'op_name', 'error_code')
_gauge_conversion_latency = monitoring.IntGauge(
'/tensorflow/lite/convert/latency', 'Conversion latency in ms.')
class TFLiteMetrics(metrics_interface.TFLiteMetricsInterface):
"""TFLite metrics helper for prod (borg) environment.
Attributes:
model_hash: A string containing the hash of the model binary.
model_path: A string containing the path of the model for debugging
purposes.
"""
def __init__(self,
model_hash: Optional[Text] = None,
model_path: Optional[Text] = None) -> None:
del self # Temporarily removing self until parameter logic is implemented.
if model_hash and not model_path or not model_hash and model_path:
raise ValueError('Both model metadata(model_hash, model_path) should be '
'given at the same time.')
if model_hash:
# TODO(b/180400857): Create stub once the service is implemented.
pass
def increase_counter_debugger_creation(self):
_counter_debugger_creation.get_cell().increase_by(1)
def increase_counter_interpreter_creation(self):
_counter_interpreter_creation.get_cell('python').increase_by(1)
def increase_counter_converter_attempt(self):
_counter_conversion_attempt.get_cell().increase_by(1)
def increase_counter_converter_success(self):
_counter_conversion_success.get_cell().increase_by(1)
def set_converter_param(self, name, value):
_gauge_conversion_params.get_cell(name).set(value)
def set_converter_error(
self, error_data: converter_error_data_pb2.ConverterErrorData):
error_code_str = converter_error_data_pb2.ConverterErrorData.ErrorCode.Name(
error_data.error_code)
_gauge_conversion_errors.get_cell(
error_data.component,
error_data.subcomponent,
error_data.operator.name,
error_code_str,
).set(error_data.error_message)
def set_converter_latency(self, value):
_gauge_conversion_latency.get_cell().set(value)
class TFLiteConverterMetrics(TFLiteMetrics):
"""Similar to TFLiteMetrics but specialized for converter.
A unique session id will be created for each new TFLiteConverterMetrics.
"""
def __init__(self) -> None:
super(TFLiteConverterMetrics, self).__init__()
session_id = uuid.uuid4().hex
self._metrics_exporter = metrics_wrapper.MetricsWrapper(session_id)
self._exported = False
def __del__(self):
if not self._exported:
self.export_metrics()
def set_export_required(self):
self._exported = False
def export_metrics(self):
self._metrics_exporter.ExportMetrics()
self._exported = True
tensorflow/lite/python/metrics/metrics_nonportable_test.py
+570
View File
@@ -0,0 +1,570 @@
# Copyright 2021 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 metrics helper TFLiteMetrics check."""
import gc
import os
import tempfile
import time
from unittest import mock
from absl.testing import parameterized
import numpy as np
import tensorflow as tf
from tensorflow.compiler.mlir.lite.metrics import converter_error_data_pb2
from tensorflow.core.framework import graph_pb2
from tensorflow.lite.python import lite
from tensorflow.lite.python.convert import ConverterError
from tensorflow.lite.python.convert import register_custom_opdefs
from tensorflow.lite.python.metrics import metrics
from tensorflow.python.client import session
from tensorflow.python.eager import context
from tensorflow.python.eager import monitoring
from tensorflow.python.framework import convert_to_constants
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import test_util
from tensorflow.python.framework.importer import import_graph_def
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.ops import string_ops
from tensorflow.python.ops.ragged import ragged_tensor
from tensorflow.python.platform import resource_loader
from tensorflow.python.platform import test
from tensorflow.python.saved_model import saved_model
from tensorflow.python.trackable import autotrackable
class MetricsNonportableTest(test_util.TensorFlowTestCase):
def test_TFLiteMetrics_creation_no_arg_success(self):
metrics.TFLiteMetrics()
def test_TFLiteMetrics_creation_arg_success(self):
metrics.TFLiteMetrics('hash', '/path/to/model')
def test_TFLiteMetrics_creation_fails_with_only_hash(self):
with self.assertRaises(ValueError):
metrics.TFLiteMetrics(model_hash='hash')
def test_TFLiteMetrics_creation_fail2_with_only_model_path(self):
with self.assertRaises(ValueError):
metrics.TFLiteMetrics(model_path='/path/to/model')
def test_debugger_creation_counter_increase_multiple_same_topic_success(self):
try:
stub = metrics.TFLiteMetrics()
stub.increase_counter_debugger_creation()
self.assertEqual(metrics._counter_debugger_creation.get_cell().value(), 1)
stub2 = metrics.TFLiteMetrics()
stub2.increase_counter_debugger_creation()
self.assertEqual(metrics._counter_debugger_creation.get_cell().value(), 2)
del stub
gc.collect()
stub2.increase_counter_debugger_creation()
self.assertEqual(metrics._counter_debugger_creation.get_cell().value(), 3)
except:
raise Exception('No exception should be raised.')
def test_interpreter_creation_counter_increase_success(self):
stub = metrics.TFLiteMetrics()
stub.increase_counter_interpreter_creation()
self.assertEqual(
metrics._counter_interpreter_creation.get_cell('python').value(), 1)
def test_converter_attempt_counter_increase_success(self):
stub = metrics.TFLiteMetrics()
stub.increase_counter_converter_attempt()
self.assertEqual(metrics._counter_conversion_attempt.get_cell().value(), 1)
def test_converter_success_counter_increase_success(self):
stub = metrics.TFLiteMetrics()
stub.increase_counter_converter_success()
self.assertEqual(metrics._counter_conversion_success.get_cell().value(), 1)
def test_converter_params_set_success(self):
stub = metrics.TFLiteMetrics()
stub.set_converter_param('name', 'value')
self.assertEqual(
metrics._gauge_conversion_params.get_cell('name').value(), 'value')
def test_converter_params_multiple_set_success(self):
stub = metrics.TFLiteMetrics()
stub.set_converter_param('name', 'value')
stub.set_converter_param('name', 'value1')
self.assertEqual(
metrics._gauge_conversion_params.get_cell('name').value(), 'value1')
def test_converter_params_multiple_label_success(self):
stub = metrics.TFLiteMetrics()
stub.set_converter_param('name1', 'value1')
stub.set_converter_param('name2', 'value2')
self.assertEqual(
metrics._gauge_conversion_params.get_cell('name1').value(), 'value1')
self.assertEqual(
metrics._gauge_conversion_params.get_cell('name2').value(), 'value2')
def test_converter_params_set_latency(self):
stub = metrics.TFLiteMetrics()
stub.set_converter_latency(34566)
self.assertEqual(metrics._gauge_conversion_latency.get_cell().value(),
34566)
class ConverterMetricsTest(test_util.TensorFlowTestCase):
"""Testing conversion metrics."""
def _constructGraphDef(self):
with ops.Graph().as_default():
in_tensor = array_ops.placeholder(
shape=[None, 16, 16, 3], dtype=dtypes.float32, name='in_tensor')
math_ops.add(in_tensor, in_tensor, name='add')
sess = session.Session()
return (
convert_to_constants.convert_variables_to_constants_from_session_graph(
sess, sess.graph_def, ['add']))
def test_conversion_from_constructor_success(self):
frozen_graph_def = self._constructGraphDef()
# Check metrics when conversion successed.
converter = lite.TFLiteConverter(frozen_graph_def, None, None,
[('in_tensor', [2, 16, 16, 3])], ['add'])
mock_metrics = mock.create_autospec(
metrics.TFLiteConverterMetrics, instance=True)
converter._tflite_metrics = mock_metrics
tflite_model = converter.convert()
self.assertIsNotNone(tflite_model)
mock_metrics.assert_has_calls([
mock.call.increase_counter_converter_attempt(),
mock.call.increase_counter_converter_success(),
mock.call.export_metrics(),
mock.call.set_converter_param('input_format', '1'),
mock.call.set_converter_param('allow_custom_ops', 'False'),
mock.call.set_converter_param('api_version', '1'),
], any_order=True) # pyformat: disable
def test_conversion_from_constructor_fail(self):
frozen_graph_def = self._constructGraphDef()
# Check metrics when conversion failed.
converter = lite.TFLiteConverter(frozen_graph_def, None, None,
[('wrong_tensor', [2, 16, 16, 3])],
['add'])
mock_metrics = mock.create_autospec(
metrics.TFLiteConverterMetrics, instance=True)
converter._tflite_metrics = mock_metrics
with self.assertRaises(ConverterError):
converter.convert()
mock_metrics.assert_has_calls([
mock.call.increase_counter_converter_attempt(),
mock.call.set_converter_param('output_format', '2'),
mock.call.set_converter_param('select_user_tf_ops', 'None'),
mock.call.set_converter_param('post_training_quantize', 'False'),
], any_order=True) # pyformat: disable
mock_metrics.increase_counter_converter_success.assert_not_called()
def _getIntegerQuantizeModel(self):
np.random.seed(0)
root = autotrackable.AutoTrackable()
@tf.function(
input_signature=[tf.TensorSpec(shape=[1, 5, 5, 3], dtype=tf.float32)])
def func(inp):
conv = tf.nn.conv2d(
inp, tf.ones([3, 3, 3, 16]), strides=[1, 1, 1, 1], padding='SAME')
output = tf.nn.relu(conv, name='output')
return output
def calibration_gen():
for _ in range(5):
yield [np.random.uniform(-1, 1, size=(1, 5, 5, 3)).astype(np.float32)]
root.f = func
to_save = root.f.get_concrete_function()
return (root, to_save, calibration_gen)
def test_conversion_from_frozen_graph_v2(self):
model, func, calibration_gen = self._getIntegerQuantizeModel()
quantized_converter = lite.TFLiteConverterV2.from_concrete_functions([func],
model)
mock_metrics = mock.create_autospec(
metrics.TFLiteConverterMetrics, instance=True)
quantized_converter._tflite_metrics = mock_metrics
quantized_converter.optimizations = [lite.Optimize.DEFAULT]
quantized_converter.representative_dataset = calibration_gen
quantized_tflite_model = quantized_converter.convert()
self.assertIsNotNone(quantized_tflite_model)
mock_metrics.assert_has_calls([
mock.call.increase_counter_converter_attempt(),
mock.call.increase_counter_converter_success(),
mock.call.set_converter_param(
'optimization_post_training_integer_quantize', 'True'),
mock.call.set_converter_param('inference_type', 'tf.int8'),
mock.call.set_converter_param('select_user_tf_ops', 'None'),
mock.call.set_converter_param('activations_type', 'tf.int8'),
], any_order=True) # pyformat: disable
def test_conversion_from_keras_v2(self):
x = [-1, 0, 1, 2, 3, 4]
y = [-3, -1, 1, 3, 5, 7]
model = tf.keras.models.Sequential(
[tf.keras.layers.Dense(units=1, input_shape=[1])])
model.compile(optimizer='sgd', loss='mean_squared_error')
model.fit(x, y, epochs=1)
converter = lite.TFLiteConverterV2.from_keras_model(model)
mock_metrics = mock.create_autospec(
metrics.TFLiteConverterMetrics, instance=True)
converter._tflite_metrics = mock_metrics
converter.convert()
mock_metrics.assert_has_calls([
mock.call.increase_counter_converter_attempt(),
mock.call.increase_counter_converter_success(),
mock.call.export_metrics(),
mock.call.set_converter_param('inference_type', 'tf.float32'),
mock.call.set_converter_param('target_ops', 'TFLITE_BUILTINS'),
mock.call.set_converter_param('optimization_default', 'False'),
], any_order=True) # pyformat: disable
def _createV1SavedModel(self, shape):
"""Create a simple SavedModel."""
saved_model_dir = os.path.join(self.get_temp_dir(), 'simple_savedmodel')
with tf.Graph().as_default():
with tf.compat.v1.Session() as sess:
in_tensor_1 = tf.compat.v1.placeholder(
shape=shape, dtype=tf.float32, name='inputB')
in_tensor_2 = tf.compat.v1.placeholder(
shape=shape, dtype=tf.float32, name='inputA')
variable_node = tf.Variable(1.0, name='variable_node')
out_tensor = in_tensor_1 + in_tensor_2 * variable_node
inputs = {'x': in_tensor_1, 'y': in_tensor_2}
outputs = {'z': out_tensor}
sess.run(tf.compat.v1.variables_initializer([variable_node]))
saved_model.simple_save(sess, saved_model_dir, inputs, outputs)
return saved_model_dir
def test_conversion_from_saved_model(self):
saved_model_dir = self._createV1SavedModel(shape=[1, 16, 16, 3])
converter = lite.TFLiteSavedModelConverter(saved_model_dir, set(['serve']),
['serving_default'])
converter.experimental_new_converter = True
mock_metrics = mock.create_autospec(
metrics.TFLiteConverterMetrics, instance=True)
converter._tflite_metrics = mock_metrics
time.process_time = mock.Mock(side_effect=np.arange(1, 1000, 2).tolist())
converter.convert()
mock_metrics.assert_has_calls([
mock.call.increase_counter_converter_attempt(),
mock.call.increase_counter_converter_success(),
mock.call.set_converter_latency(2000),
mock.call.export_metrics(),
], any_order=True) # pyformat: disable
def disable_converter_counter_metrics(self, tflite_metrics):
def empty_func():
pass
tflite_metrics.increase_counter_converter_attempt = empty_func
tflite_metrics.increase_counter_converter_success = empty_func
def test_export_at_conversion_done(self):
saved_model_dir = self._createV1SavedModel(shape=[1, 16, 16, 3])
converter = lite.TFLiteConverterV2.from_saved_model(saved_model_dir)
tflite_metrics = converter._tflite_metrics
mock_exporter = mock.MagicMock()
tflite_metrics._metrics_exporter = mock_exporter
self.disable_converter_counter_metrics(tflite_metrics)
mock_exporter.ExportMetrics.assert_not_called()
converter.convert()
mock_exporter.ExportMetrics.assert_called_once()
tflite_metrics.__del__()
mock_exporter.ExportMetrics.assert_called_once()
def test_export_at_exit(self):
saved_model_dir = self._createV1SavedModel(shape=[1, 16, 16, 3])
converter = lite.TFLiteConverterV2.from_saved_model(saved_model_dir)
tflite_metrics = converter._tflite_metrics
mock_exporter = mock.MagicMock()
tflite_metrics._metrics_exporter = mock_exporter
self.disable_converter_counter_metrics(tflite_metrics)
mock_exporter.ExportMetrics.assert_not_called()
tflite_metrics.__del__()
mock_exporter.ExportMetrics.assert_called_once()
def mock_ngrams(data, width, axis=-1, string_separator=' ', name=None):
"""This mock Ngrams lack the width attr, causing conversion to fail."""
experimental_implements = [
'name: "tftext:Ngrams"',
'attr { key: "axis" value { i: %d } }' % axis,
'attr { key: "reduction_type" value { s: "STRING_JOIN" } }',
'attr { key: "string_separator" value { s: "%s" } }' % string_separator,
]
experimental_implements = ' '.join(experimental_implements)
@tf.function(experimental_implements=experimental_implements)
def func(data):
with ops.name_scope(name, 'NGrams', [data, width]):
data = ragged_tensor.convert_to_tensor_or_ragged_tensor(data, name='data')
slices = []
for start in range(width):
stop = None if start - width + 1 == 0 else start - width + 1
if axis >= 0:
idx = [slice(None)] * axis + [slice(start, stop)]
else:
idx = [Ellipsis, slice(start, stop)] + [slice(None)] * (-axis - 1)
slices.append(data[idx])
# Stack the slices.
stack_axis = axis + 1 if axis >= 0 else axis
windowed_data = array_ops_stack.stack(slices, stack_axis)
return string_ops.reduce_join(
windowed_data, axis=axis, separator=string_separator)
return func(data)
class ConverterErrorMetricTest(test_util.TensorFlowTestCase,
parameterized.TestCase):
"""Testing conversion error metric."""
def setUp(self):
super(ConverterErrorMetricTest, self).setUp()
# Mock metrics instance except errors so other test cases are not affected.
mock_attempt = mock.create_autospec(monitoring.Counter, instance=True)
self._counter_conversion_attempt = metrics._counter_conversion_attempt
metrics._counter_conversion_attempt = mock_attempt
mock_success = mock.create_autospec(monitoring.Counter, instance=True)
self._counter_conversion_success = metrics._counter_conversion_success
metrics._counter_conversion_success = mock_success
mock_params = mock.create_autospec(monitoring.StringGauge, instance=True)
self._gauge_conversion_params = metrics._gauge_conversion_params
metrics._gauge_conversion_params = mock_params
def tearDown(self):
super(ConverterErrorMetricTest, self).tearDown()
# # Restore metrics instances.
metrics._counter_conversion_attempt = self._counter_conversion_attempt
metrics._counter_conversion_success = self._counter_conversion_success
metrics._gauge_conversion_params = self._gauge_conversion_params
def convert_and_check_location_info(self,
converter,
expected_type,
expected_sources=None):
# The custom attribute of ConverterError can't be accessed with
# assertRaises so use try-catch block instead.
try:
tflite_model = converter.convert()
self.assertIsNone(tflite_model)
except ConverterError as converter_error:
# pylint: disable=g-assert-in-except
self.assertLen(converter_error.errors, 1)
location = converter_error.errors[0].location
self.assertEqual(location.type, expected_type)
if expected_sources:
debug_string = str(location)
for source in expected_sources:
self.assertIn(source, debug_string)
# pylint: enable=g-assert-in-except
def test_failure_at_PrepareCompositeFunctionsPass(self):
if context.is_tfrt_enabled():
self.skipTest('This test crashed with TFRT.')
class NgramsLayer(tf.keras.layers.Layer):
def call(self, input_tensor, **kwargs):
return mock_ngrams(input_tensor, width=2, axis=-1, string_separator=' ')
# Registers a fake WhitespaceTokenizeWithOffsets so the TFText fusing logic
# is enable in MLIR side.
custom_opdefs_str = (
'name: \'WhitespaceTokenizeWithOffsets\' input_arg: {name: \'Input1\' '
'type: DT_FLOAT} input_arg: {name: \'Input2\' type: DT_FLOAT} '
'output_arg: {name: \'Output\' type: DT_FLOAT}')
register_custom_opdefs([custom_opdefs_str])
model = tf.keras.models.Sequential([NgramsLayer()])
model.predict(tf.constant(['test']))
converter = lite.TFLiteConverterV2.from_keras_model(model)
converter.allow_custom_ops = True
self.convert_and_check_location_info(
converter, converter_error_data_pb2.ConverterErrorData.UNKNOWNLOC)
exported_error = metrics._gauge_conversion_errors.get_cell(
'CONVERT_TF_TO_TFLITE_MODEL',
'PrepareCompositeFunctionsPass',
'tf.Const',
'UNKNOWN',
).value()
self.assertEqual(exported_error,
"\'width\' attribute is not set or not an integer")
def test_need_flex_ops(self):
def create_graph_with_custom_add(opname='CustomAdd'):
custom_opdefs_str = (
'name: \'' + opname +
'\' input_arg: {name: \'Input1\' type: DT_FLOAT} '
'input_arg: {name: \'Input2\' type: DT_FLOAT} output_arg: {name: '
'\'Output\' type: DT_FLOAT}')
# Create a graph that has one add op.
new_graph = graph_pb2.GraphDef()
with ops.Graph().as_default():
with session.Session() as sess:
in_tensor = array_ops.placeholder(
shape=[1, 16, 16, 3], dtype=dtypes.float32, name='input')
out_tensor = in_tensor + in_tensor
inputs = {'x': in_tensor}
outputs = {'z': out_tensor}
new_graph.CopyFrom(sess.graph_def)
# Rename Add op name to opname.
for node in new_graph.node:
if node.op.startswith('Add'):
node.op = opname
del node.attr['T']
# Register custom op defs to import modified graph def.
register_custom_opdefs([custom_opdefs_str])
return (new_graph, inputs, outputs)
new_graph, inputs, outputs = create_graph_with_custom_add()
# 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)
self.convert_and_check_location_info(
converter,
converter_error_data_pb2.ConverterErrorData.NAMELOC,
expected_sources='add')
exported_error = metrics._gauge_conversion_errors.get_cell(
'CONVERT_TF_TO_TFLITE_MODEL', 'CONVERT_SAVED_MODEL', 'tf.CustomAdd',
'ERROR_NEEDS_CUSTOM_OPS').value()
self.assertIn(
"'tf.CustomAdd' op is neither a custom op nor a flex op\n",
exported_error,
)
self.assertIn('Error code: ERROR_NEEDS_CUSTOM_OPS', exported_error)
def test_unsupported_control_flow_v1(self):
filename = resource_loader.get_path_to_datafile(
'../testdata/control_flow_v1_saved_model')
converter = lite.TFLiteConverterV2.from_saved_model(filename)
self.convert_and_check_location_info(
converter, converter_error_data_pb2.ConverterErrorData.UNKNOWNLOC)
exported_error = metrics._gauge_conversion_errors.get_cell(
'CONVERT_TF_TO_TFLITE_MODEL', 'CONVERT_SAVED_MODEL', '',
'ERROR_UNSUPPORTED_CONTROL_FLOW_V1').value()
self.assertEqual(
exported_error,
'Merge only has 4 inputs, while only merge nodes with two inputs '
'supported.\n\tFailed to functionalize Control Flow V1 ops. Consider '
'using Control Flow V2 ops instead. See https://www.tensorflow.org/'
'api_docs/python/tf/compat/v1/enable_control_flow_v2.')
def test_location_from_concrete_functions(self):
@tf.function(input_signature=[
tf.TensorSpec(shape=[None, None, 2, 3, 3], dtype=tf.complex64),
tf.TensorSpec(shape=[None, None, 1, 3, 3], dtype=tf.complex64),
])
def model(a, b):
return tf.add(a, b, name='add')
converter = lite.TFLiteConverterV2.from_concrete_functions(
[model.get_concrete_function()], model)
self.convert_and_check_location_info(
converter,
converter_error_data_pb2.ConverterErrorData.CALLSITELOC,
expected_sources=[
'tensorflow/lite/python/metrics/metrics_nonportable_test.py',
])
def test_location_from_saved_model(self):
with tempfile.TemporaryDirectory() as tmp_dir:
class Adder(tf.Module):
@tf.function(input_signature=[
tf.TensorSpec(shape=[None, None, 2, 3, 3], dtype=tf.complex64),
tf.TensorSpec(shape=[None, None, 1, 3, 3], dtype=tf.complex64),
])
def serving_default(self, a, b):
return tf.add(a, b, name='add')
tf.saved_model.save(
Adder(),
tmp_dir,
options=tf.saved_model.SaveOptions(save_debug_info=True))
converter = lite.TFLiteConverterV2.from_saved_model(tmp_dir)
self.convert_and_check_location_info(
converter,
converter_error_data_pb2.ConverterErrorData.CALLSITELOC,
expected_sources=[
'tensorflow/lite/python/metrics/metrics_nonportable_test.py',
])
@parameterized.named_parameters(
('_WithoutLoweringToSavedModel', False, None),
('_WithLoweringToSavedModel', True,
'tensorflow/lite/python/metrics/metrics_nonportable_test.py'))
def test_location_from_keras_model(self, lower_to_saved_model,
expected_source):
input_tensor1 = tf.keras.layers.Input(
shape=[None, None, 2, 3, 3], dtype=tf.complex64)
input_tensor2 = tf.keras.layers.Input(
shape=[None, None, 2, 3, 3], dtype=tf.complex64)
output = tf.keras.layers.Add()([input_tensor1, input_tensor2])
model = tf.keras.Model(
inputs=[input_tensor1, input_tensor2], outputs=output)
model.compile(
optimizer='adam',
loss='sparse_categorical_crossentropy',
metrics=['accuracy'])
converter = lite.TFLiteConverterV2.from_keras_model(model)
converter.experimental_lower_to_saved_model = lower_to_saved_model
# The location does not contain callsite to the current file.
self.convert_and_check_location_info(
converter,
converter_error_data_pb2.ConverterErrorData.CALLSITELOC,
expected_sources=[expected_source] if expected_source else None)
if __name__ == '__main__':
test.main()
tensorflow/lite/python/metrics/metrics_portable.py
+70
View File
@@ -0,0 +1,70 @@
# Copyright 2021 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.
# ==============================================================================
"""Python TFLite metrics helper."""
import os
from typing import Optional, Text
# pylint: disable=g-import-not-at-top
if not os.path.splitext(__file__)[0].endswith(
os.path.join('tflite_runtime', 'metrics_portable')):
# This file is part of tensorflow package.
from tensorflow.lite.python.metrics import metrics_interface # type: ignore
else:
# This file is part of tflite_runtime package.
from tflite_runtime import metrics_interface # type: ignore
# pylint: enable=g-import-not-at-top
class TFLiteMetrics(metrics_interface.TFLiteMetricsInterface):
"""TFLite metrics helper."""
def __init__(self,
model_hash: Optional[Text] = None,
model_path: Optional[Text] = None) -> None:
pass
def increase_counter_debugger_creation(self):
pass
def increase_counter_interpreter_creation(self):
pass
def increase_counter_converter_attempt(self):
pass
def increase_counter_converter_success(self):
pass
def set_converter_param(self, name, value):
pass
def set_converter_error(self, error_data):
pass
def set_converter_latency(self, value):
pass
class TFLiteConverterMetrics(TFLiteMetrics):
"""Similar to TFLiteMetrics but specialized for converter."""
def __del__(self):
pass
def set_export_required(self):
pass
def export_metrics(self):
pass
tensorflow/lite/python/metrics/metrics_portable_test.py
+48
View File
@@ -0,0 +1,48 @@
# Copyright 2021 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 metrics helpr TFLiteMetrics check."""
from tensorflow.lite.python.metrics import metrics
from tensorflow.python.framework import test_util
from tensorflow.python.platform import test
class MetricsPortableTest(test_util.TensorFlowTestCase):
def test_TFLiteMetrics_creation_success(self):
metrics.TFLiteMetrics()
def test_debugger_creation_counter_increase_success(self):
stub = metrics.TFLiteMetrics()
stub.increase_counter_debugger_creation()
def test_interpreter_creation_counter_increase_success(self):
stub = metrics.TFLiteMetrics()
stub.increase_counter_interpreter_creation()
def test_converter_attempt_counter_increase_success(self):
stub = metrics.TFLiteMetrics()
stub.increase_counter_converter_attempt()
def test_converter_success_counter_increase_success(self):
stub = metrics.TFLiteMetrics()
stub.increase_counter_converter_success()
def test_converter_params_set_success(self):
stub = metrics.TFLiteMetrics()
stub.set_converter_param('name', 'value')
if __name__ == '__main__':
test.main()
tensorflow/lite/python/metrics/wrapper/metrics_wrapper.h
+58
View File
@@ -0,0 +1,58 @@
/* Copyright 2021 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_PYTHON_METRICS_WRAPPER_METRICS_WRAPPER_H_
#define TENSORFLOW_LITE_PYTHON_METRICS_WRAPPER_METRICS_WRAPPER_H_
#include <memory>
#include <string>
// Place `<locale>` before <Python.h> to avoid build failures in macOS.
#include <locale>
// The empty line above is on purpose as otherwise clang-format will
// automatically move <Python.h> before <locale>.
#include <Python.h>
// We forward declare TFLite classes here to avoid exposing them to SWIG.
namespace tensorflow {
namespace monitoring {
class MetricsExporter;
} // namespace monitoring
} // namespace tensorflow
namespace tflite {
namespace metrics_wrapper {
class MetricsWrapper {
public:
using MetricsExporter = tensorflow::monitoring::MetricsExporter;
// SWIG caller takes ownership of pointer.
static MetricsWrapper* CreateMetricsWrapper(const std::string& session_id);
~MetricsWrapper();
// Export metrics with Streamz.
PyObject* ExportMetrics();
private:
explicit MetricsWrapper(std::unique_ptr<MetricsExporter> exporter);
const std::unique_ptr<MetricsExporter> exporter_;
};
} // namespace metrics_wrapper
} // namespace tflite
#endif // TENSORFLOW_LITE_PYTHON_METRICS_WRAPPER_METRICS_WRAPPER_H_
tensorflow/lite/python/metrics/wrapper/metrics_wrapper.py
+34
View File
@@ -0,0 +1,34 @@
# Copyright 2021 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.
# ==============================================================================
"""Stub to make pywrap metrics wrapper accessible."""
from tensorflow.compiler.mlir.lite.metrics import converter_error_data_pb2
from tensorflow.compiler.mlir.lite.python import wrap_converter
from tensorflow.lite.python.metrics._pywrap_tensorflow_lite_metrics_wrapper import MetricsWrapper # pylint: disable=unused-import
def retrieve_collected_errors():
"""Returns and clears the list of collected errors in ErrorCollector.
The RetrieveCollectedErrors function in C++ returns a list of serialized proto
messages. This function will convert them to ConverterErrorData instances.
Returns:
A list of ConverterErrorData.
"""
serialized_message_list = wrap_converter.wrapped_retrieve_collected_errors()
return list(
map(converter_error_data_pb2.ConverterErrorData.FromString,
serialized_message_list))
tensorflow/lite/python/metrics/wrapper/metrics_wrapper_nonportable.cc
+63
View File
@@ -0,0 +1,63 @@
/* Copyright 2021 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.
==============================================================================*/
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "learning/brain/google/monitoring/metrics_exporter.h"
#include "tensorflow/lite/python/metrics/wrapper/metrics_wrapper.h"
namespace tflite {
namespace metrics_wrapper {
MetricsWrapper::MetricsWrapper(std::unique_ptr<MetricsExporter> exporter)
: exporter_(std::move(exporter)) {}
MetricsWrapper::~MetricsWrapper() {}
MetricsWrapper* MetricsWrapper::CreateMetricsWrapper(
const std::string& session_id) {
MetricsExporter::Options options;
options.export_at_exit = false;
{
// Add session_id to root label list.
std::vector<streamz::Entity::Label> root_labels;
streamz::Entity::Label session_id_label;
session_id_label.set_key("session_id");
session_id_label.set_string_value(session_id);
root_labels.push_back(session_id_label);
options.entity_labels = root_labels;
}
std::unique_ptr<MetricsExporter> exporter(new MetricsExporter(options));
MetricsWrapper* wrapper = new MetricsWrapper(std::move(exporter));
return wrapper;
}
PyObject* MetricsWrapper::ExportMetrics() {
if (!exporter_) {
PyErr_SetString(PyExc_ValueError, "MetricsExporter was not initialized.");
return nullptr;
}
exporter_->ExportMetrics();
Py_RETURN_NONE;
}
} // namespace metrics_wrapper
} // namespace tflite
tensorflow/lite/python/metrics/wrapper/metrics_wrapper_portable.cc
+60
View File
@@ -0,0 +1,60 @@
/* Copyright 2021 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.
==============================================================================*/
#include <memory>
#include <utility>
#include <vector>
#include "tensorflow/lite/python/metrics/wrapper/metrics_wrapper.h"
namespace tensorflow {
namespace monitoring {
class MetricsExporter {
public:
MetricsExporter() {}
void ExportMetrics() {}
};
} // namespace monitoring
} // namespace tensorflow
namespace tflite {
namespace metrics_wrapper {
MetricsWrapper::MetricsWrapper(std::unique_ptr<MetricsExporter> exporter)
: exporter_(std::move(exporter)) {}
MetricsWrapper::~MetricsWrapper() {}
MetricsWrapper* MetricsWrapper::CreateMetricsWrapper(
const std::string& session_id) {
std::unique_ptr<MetricsExporter> exporter(new MetricsExporter());
MetricsWrapper* wrapper = new MetricsWrapper(std::move(exporter));
return wrapper;
}
PyObject* MetricsWrapper::ExportMetrics() {
if (!exporter_) {
PyErr_SetString(PyExc_ValueError, "MetricsExporter was not initialized.");
return nullptr;
}
exporter_->ExportMetrics();
Py_RETURN_NONE;
}
} // namespace metrics_wrapper
} // namespace tflite
tensorflow/lite/python/metrics/wrapper/metrics_wrapper_pybind11.cc
+44
View File
@@ -0,0 +1,44 @@
/* Copyright 2021 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.
==============================================================================*/
#include <string>
#include "pybind11/pybind11.h" // from @pybind11
#include "pybind11/pytypes.h" // from @pybind11
#include "pybind11/stl.h" // from @pybind11
#include "tensorflow/lite/python/metrics/wrapper/metrics_wrapper.h"
#include "tensorflow/python/lib/core/pybind11_lib.h"
namespace py = pybind11;
using tflite::metrics_wrapper::MetricsWrapper;
PYBIND11_MODULE(_pywrap_tensorflow_lite_metrics_wrapper, m) {
m.doc() = R"pbdoc(
_pywrap_tensorflow_lite_metrics_wrapper
-----
)pbdoc";
py::class_<MetricsWrapper>(m, "MetricsWrapper")
.def(py::init([](const std::string& session_id) {
auto* wrapper = MetricsWrapper::CreateMetricsWrapper(session_id);
if (!wrapper) {
throw std::invalid_argument("Failed to created MetricsWrapper");
}
return wrapper;
}))
.def("ExportMetrics", [](MetricsWrapper& self) {
return tensorflow::PyoOrThrow(self.ExportMetrics());
});
}
tensorflow/lite/python/metrics/wrapper/metrics_wrapper_test.py
+50
View File
@@ -0,0 +1,50 @@
# Copyright 2021 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.
# ==============================================================================
"""TFLite metrics_wrapper module test cases."""
import tensorflow as tf
from tensorflow.lite.python import lite
from tensorflow.lite.python.convert import ConverterError
from tensorflow.lite.python.metrics.wrapper import metrics_wrapper
from tensorflow.python.framework import test_util
from tensorflow.python.platform import test
class MetricsWrapperTest(test_util.TensorFlowTestCase):
def test_basic_retrieve_collected_errors_empty(self):
errors = metrics_wrapper.retrieve_collected_errors()
self.assertEmpty(errors)
def test_basic_retrieve_collected_errors_not_empty(self):
@tf.function(
input_signature=[tf.TensorSpec(shape=[None], dtype=tf.float32)])
def func(x):
return tf.cosh(x)
converter = lite.TFLiteConverterV2.from_concrete_functions(
[func.get_concrete_function()], func)
try:
converter.convert()
except ConverterError as err:
# retrieve_collected_errors is already captured in err.errors
captured_errors = err.errors
self.assertNotEmpty(captured_errors)
if __name__ == "__main__":
test.main()
tensorflow/lite/python/op_hint.py
+1338
View File
File diff suppressed because it is too large Load Diff
tensorflow/lite/python/optimize/BUILD
+109
View File
@@ -0,0 +1,109 @@
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
load("//tensorflow:tensorflow.bzl", "py_test")
load("//tensorflow:tensorflow.default.bzl", "pybind_extension")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
default_visibility = ["//visibility:public"],
licenses = ["notice"],
)
cc_library(
name = "calibration_wrapper_lib",
srcs = ["calibration_wrapper.cc"],
hdrs = ["calibration_wrapper.h"],
deps = [
"//tensorflow/compiler/mlir/lite:offset_buffer",
"//tensorflow/compiler/mlir/lite/schema:schema_fbs_with_mutable",
"//tensorflow/lite:framework",
"//tensorflow/lite:shared_library",
"//tensorflow/lite:string_util",
"//tensorflow/lite/core:framework",
"//tensorflow/lite/core/c:common",
"//tensorflow/lite/core/kernels:builtin_ops",
"//tensorflow/lite/python/interpreter_wrapper:numpy",
"//tensorflow/lite/python/interpreter_wrapper:python_error_reporter",
"//tensorflow/lite/python/interpreter_wrapper:python_utils",
"//tensorflow/lite/tools/optimize:quantization_wrapper_utils",
"//tensorflow/lite/tools/optimize:quantize_model",
"//tensorflow/lite/tools/optimize/calibration:calibration_reader",
"//tensorflow/lite/tools/optimize/calibration:calibrator_lib",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:str_format",
"@com_google_absl//absl/types:optional",
"@xla//third_party/python_runtime:headers", # buildcleaner: keep
],
)
pybind_extension(
name = "_pywrap_tensorflow_lite_calibration_wrapper",
srcs = [
"calibration_wrapper_pybind11.cc",
],
hdrs = ["calibration_wrapper.h"],
additional_stubgen_deps = [
"//third_party/py/numpy:numpy",
],
common_lib_packages = [
"litert/python",
"tensorflow/lite/python",
],
enable_stub_generation = True,
link_in_framework = True,
pytype_srcs = [
"_pywrap_tensorflow_lite_calibration_wrapper.pyi",
],
wrap_py_init = True,
deps = [
":calibration_wrapper_lib",
"//tensorflow/lite:framework",
"//tensorflow/lite/core:framework_stable",
"//tensorflow/python/lib/core:pybind11_lib",
"@pybind11",
"@xla//third_party/python_runtime:headers",
],
)
py_library(
name = "calibrator",
srcs = [
"calibrator.py",
],
strict_deps = True,
visibility = ["//visibility:public"],
deps = [
":_pywrap_tensorflow_lite_calibration_wrapper", # buildcleaner: keep
"//tensorflow/lite/python:convert_phase",
"//tensorflow/lite/python:interpreter",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/util:lazy_loader",
"//third_party/py/numpy",
],
)
py_test(
name = "calibrator_test",
srcs = ["calibrator_test.py"],
data = [
":test_data",
"//tensorflow/lite:testdata/multi_add.bin",
],
strict_deps = True,
tags = ["no_oss"],
deps = [
":calibrator",
"@absl_py//absl/testing:parameterized",
#internal proto upb dep
"//third_party/py/numpy",
"//tensorflow:tensorflow_py_no_contrib",
"//tensorflow/lite/python:lite",
"//tensorflow/lite/python:schema_py",
"//tensorflow/lite/tools:flatbuffer_utils",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/platform:resource_loader",
],
)
tensorflow/lite/python/optimize/_pywrap_tensorflow_lite_calibration_wrapper.pyi
+38
View File
@@ -0,0 +1,38 @@
# Copyright 2023 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.
# ==============================================================================
from typing import Callable, overload
class CalibrationWrapper:
def __init__(self, arg0: object, arg1: list[str], arg2: list[Callable[[int], None]]) -> None: ...
def Calibrate(self) -> object: ...
@overload
def FeedTensor(self, arg0: object, arg1: str) -> object: ...
@overload
def FeedTensor(self, arg0: object) -> object: ...
@overload
def Prepare(self, arg0: object, arg1: str) -> object: ...
@overload
def Prepare(self, arg0: object) -> object: ...
@overload
def Prepare(self, arg0: str) -> object: ...
@overload
def Prepare(self) -> object: ...
@overload
def QuantizeModel(self, arg0: int, arg1: int, arg2: bool, arg3: int, arg4: int, arg5: bool, arg6: bool) -> object: ...
@overload
def QuantizeModel(self, arg0: int, arg1: int, arg2: bool, arg3: str) -> object: ...
def AddIntermediateTensors(arg0: object) -> object: ...
tensorflow/lite/python/optimize/calibration_wrapper.cc
+844
View File
@@ -0,0 +1,844 @@
/* 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.
==============================================================================*/
// clang-format off
// This #include needs to precede the inclusion of any other TF Lite header
// file that might depend on the non-mutable schema_generated.h, directly,
// e.g. core/api/op_resolver.h, or indirectly, e.g. core/subgraph.h.
// That's because "tensorflow/lite/schema/mutable/schema_generated.h"
// and "tensorflow/lite/schema/schema_generated.h" both use the same
// header guard macro (FLATBUFFERS_GENERATED_SCHEMA_TFLITE_H_), but have
// different contents (the former is a superset of the latter). In particular
// the one in mutable/ is built with the "--gen-mutable" and "--gen-object-api"
// flags to the flatbuffer schema compiler which cause some additional
// (non-virtual) accessor methods and API functions to be declared.
// The code here uses those methods, so we need to make sure that we get
// the mutable variant of this header.
#include "tensorflow/compiler/mlir/lite/schema/mutable/schema_generated.h"
#include "tensorflow/lite/python/optimize/calibration_wrapper.h"
// clang-format on
// NOLINTBEGIN
// Nolint disables warnings about header file ordering caused by
// `mutable/schema_generated.h`.
#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <functional>
#include <limits>
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <vector>
// NOLINTEND
#include "absl/algorithm/container.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_format.h"
#include "absl/strings/string_view.h"
#include "tensorflow/compiler/mlir/lite/offset_buffer.h"
#include "tensorflow/lite/core/c/common.h"
#include "tensorflow/lite/core/interpreter.h"
#include "tensorflow/lite/core/kernels/register.h"
#include "tensorflow/lite/core/model_builder.h"
#include "tensorflow/lite/mutable_op_resolver.h"
#include "tensorflow/lite/python/interpreter_wrapper/numpy.h"
#include "tensorflow/lite/python/interpreter_wrapper/python_error_reporter.h"
#include "tensorflow/lite/python/interpreter_wrapper/python_utils.h"
#include "tensorflow/lite/shared_library.h"
#include "tensorflow/lite/string_util.h"
#include "tensorflow/lite/tools/optimize/calibration/calibration_reader.h"
#include "tensorflow/lite/tools/optimize/calibration/calibrator.h"
#include "tensorflow/lite/tools/optimize/quantization_wrapper_utils.h"
#include "tensorflow/lite/tools/optimize/quantize_model.h"
#define TFLITE_PY_CHECK(x) \
if ((x) != kTfLiteOk) { \
return error_reporter_->exception(); \
}
#define TFLITE_PY_ENSURE_VALID_INTERPRETER() \
if (!interpreter_) { \
PyErr_SetString(PyExc_ValueError, "Interpreter was not initialized."); \
return nullptr; \
}
namespace tflite {
namespace calibration_wrapper {
namespace {
using ::tflite::interpreter_wrapper::PythonErrorReporter;
using ::tflite::ops::builtin::BuiltinOpResolver;
using ::tflite::optimize::AddIntermediateTensorsToFusedOp;
using ::tflite::optimize::QuantizeModelAllOperators;
using ::tflite::optimize::calibration::BuildLoggingInterpreter;
using ::tflite::optimize::calibration::CalibrationReader;
using ::tflite::python::ImportNumpy;
using ::tflite::python_utils::ConvertFromPyString;
using ::tflite::python_utils::ConvertToPyString;
using ::tflite::python_utils::PyDecrefDeleter;
using ::tflite::python_utils::TfLiteTypeFromPyArray;
using ::tflite::python_utils::TfLiteTypeFromPyType;
std::unique_ptr<ModelT> CreateMutableModel(const Model& model) {
auto copied_model = std::make_unique<ModelT>();
model.UnPackTo(copied_model.get(), nullptr);
return copied_model;
}
bool NoOpModel(const FlatBufferModel& model) {
return model->subgraphs()->size() == 1 &&
(!model->subgraphs()->begin()->operators() ||
model->subgraphs()->begin()->operators()->size() == 0);
}
inline TensorType TfLiteTypeToSchemaType(TfLiteType type) {
switch (type) {
case kTfLiteNoType:
return TensorType_FLOAT32; // TODO(b/129336260): No schema type for none.
case kTfLiteFloat32:
return TensorType_FLOAT32;
case kTfLiteFloat16:
return TensorType_FLOAT16;
case kTfLiteBFloat16:
return TensorType_BFLOAT16;
case kTfLiteFloat64:
return TensorType_FLOAT64;
case kTfLiteInt32:
return TensorType_INT32;
case kTfLiteUInt32:
return TensorType_UINT32;
case kTfLiteInt2:
return TensorType_INT2;
case kTfLiteUInt4:
return TensorType_UINT4;
case kTfLiteFloat8E4M3FN:
return TensorType_FLOAT8_E4M3FN;
case kTfLiteFloat8E5M2:
return TensorType_FLOAT8_E5M2;
case kTfLiteInt4:
return TensorType_INT4;
case kTfLiteUInt8:
return TensorType_UINT8;
case kTfLiteInt8:
return TensorType_INT8;
case kTfLiteInt64:
return TensorType_INT64;
case kTfLiteUInt64:
return TensorType_UINT64;
case kTfLiteString:
return TensorType_STRING;
case kTfLiteBool:
return TensorType_BOOL;
case kTfLiteInt16:
return TensorType_INT16;
case kTfLiteUInt16:
return TensorType_UINT16;
case kTfLiteComplex64:
return TensorType_COMPLEX64;
case kTfLiteComplex128:
return TensorType_COMPLEX128;
case kTfLiteResource:
return TensorType_RESOURCE;
case kTfLiteVariant:
return TensorType_VARIANT;
}
// No default to get compiler error when new type is introduced.
}
bool RegisterCustomOpByName(const char* registerer_name,
MutableOpResolver* resolver) {
// Registerer functions take a pointer to a BuiltinOpResolver as an input
// parameter and return void.
// TODO(b/137576229): We should implement this functionality in a more
// principled way.
typedef void (*RegistererFunctionType)(MutableOpResolver*);
// Look for the Registerer function by name.
RegistererFunctionType registerer = reinterpret_cast<RegistererFunctionType>(
SharedLibrary::GetSymbol(registerer_name));
// Fail in an informative way if the function was not found.
if (registerer == nullptr) {
PyErr_Format(PyExc_ValueError,
"Looking up symbol '%s' failed with error '%s'.",
registerer_name, SharedLibrary::GetError());
return false;
}
// Call the registerer with the resolver.
registerer(resolver);
return true;
}
// Returns the dimension from the stored list in the PyObject. If the given
// PyObject is not a list, it will return absl::optional and set the Python
// error message to notify users.
std::optional<std::vector<int>> ConvertInputShapeToVector(
PyObject* input_shapes, size_t index) {
PyObject* shape = PyList_GetItem(input_shapes, index);
if (!shape || !PyList_Check(shape)) {
PyErr_Format(PyExc_ValueError,
"Invalid %ld input shape: expected to be a list.", index);
return std::nullopt;
}
size_t size = PyList_Size(shape);
std::vector<int> dims(size);
for (size_t dim_index = 0; dim_index < size; ++dim_index) {
PyObject* dim = PyList_GetItem(shape, dim_index);
dims[dim_index] = PyLong_AsLong(dim);
}
return dims;
}
// Finds the starting position of the offset buffer within `model_buffer` if the
// `model_buffer` can be split into base buffer and offset buffer. Returns
// `std::nullopt` iff offset buffer is not used or there were no buffers with
// valid offset. Assumes `model_buffer` is valid.
std::optional<int64_t> GetOffsetBufferStartPosition(
const absl::string_view model_buffer) {
const Model& model = *GetModel(model_buffer.data());
if (!FlatBufferModel::CheckBufferOutsideModel(&model)) {
// Means the offset buffer is not used, e.g.
// `_experimental_use_buffer_offset` is not set.
return std::nullopt;
}
const int64_t int64_max = std::numeric_limits<int64_t>::max();
const int64_t min_offset = absl::c_accumulate(
*model.buffers(), /*init=*/int64_max,
/*binary_op=*/[](const int64_t acc, const Buffer* buffer) -> int64_t {
const int64_t buffer_offset = buffer->offset();
return IsValidBufferOffset(buffer_offset) ? std::min(acc, buffer_offset)
: acc;
});
if (min_offset == int64_max) {
// Means there were no buffers with valid offset.
return std::nullopt;
}
return min_offset;
}
// Splits the model buffer into base buffer and offset buffer. Offset buffer may
// exist when `_experimental_use_buffer_offset` is set.
std::pair<absl::string_view, absl::string_view> SplitOffsetBuffer(
const absl::string_view model_buffer) {
const std::optional<int64_t> offset_buffer_pos =
GetOffsetBufferStartPosition(model_buffer);
if (offset_buffer_pos == std::nullopt) {
return {model_buffer, absl::string_view(model_buffer.data(), /*len=*/0)};
}
const absl::string_view base_buffer(model_buffer.data(), *offset_buffer_pos);
const int64_t offset_buffer_length = model_buffer.size() - *offset_buffer_pos;
const absl::string_view offset_buffer(
model_buffer.data() + *offset_buffer_pos, offset_buffer_length);
return {base_buffer, offset_buffer};
}
// Merges `base_buffer` with the `offset_buffer` that contains the actual tensor
// buffer data.
std::string MergeOffsetBuffer(const absl::string_view base_buffer,
const absl::string_view offset_buffer) {
return absl::StrCat(base_buffer, offset_buffer);
}
// Updates buffer offsets in `base_buffer` by `offset_diff`.
std::string UpdateBufferOffsets(const absl::string_view base_buffer,
const int64_t offset_diff) {
std::string result_buffer(base_buffer);
Model* mutable_model = GetMutableModel(result_buffer.data());
for (Buffer* buffer : *mutable_model->mutable_buffers()) {
if (const int64_t offset = buffer->offset(); IsValidBufferOffset(offset)) {
buffer->mutate_offset(offset + offset_diff);
}
}
return result_buffer;
}
} // namespace
PyObject* AddIntermediateTensors(PyObject* data) {
char* buf = nullptr;
Py_ssize_t length;
std::unique_ptr<PythonErrorReporter> error_reporter(new PythonErrorReporter);
ImportNumpy();
if (ConvertFromPyString(data, &buf, &length) == -1) {
return nullptr;
}
std::unique_ptr<FlatBufferModel> model =
FlatBufferModel::BuildFromBuffer(buf, length, error_reporter.get());
if (!model) {
PyErr_Format(PyExc_ValueError, "Invalid model");
return nullptr;
}
const auto [base_buffer, offset_buffer] =
SplitOffsetBuffer(/*model_buffer=*/absl::string_view(buf, length));
flatbuffers::FlatBufferBuilder builder;
auto tflite_model = CreateMutableModel(*model->GetModel());
if (AddIntermediateTensorsToFusedOp(&builder, tflite_model.get()) !=
kTfLiteOk) {
error_reporter->exception();
return nullptr;
}
const int64_t result_base_buffer_size = builder.GetSize();
if (result_base_buffer_size == 0) {
// When AddIntermediateTensorsToFusedOp early returns, return the model as
// it is.
return ConvertToPyString(buf, length);
}
const int64_t offset_diff =
result_base_buffer_size - static_cast<int64_t>(base_buffer.size());
const std::string updated_result_base_buffer = UpdateBufferOffsets(
/*base_buffer=*/absl::string_view(
reinterpret_cast<const char*>(builder.GetCurrentBufferPointer()),
builder.GetSize()),
offset_diff);
const std::string result_buffer =
MergeOffsetBuffer(updated_result_base_buffer, offset_buffer);
return ConvertToPyString(result_buffer.data(), result_buffer.size());
}
CalibrationWrapper::CalibrationWrapper(
std::unique_ptr<Interpreter> interpreter,
std::unique_ptr<BuiltinOpResolver> resolver,
std::unique_ptr<PythonErrorReporter> error_reporter,
std::unique_ptr<FlatBufferModel> model,
std::unique_ptr<CalibrationReader> reader,
std::unique_ptr<std::string> model_str)
: interpreter_(std::move(interpreter)),
error_reporter_(std::move(error_reporter)),
resolver_(std::move(resolver)),
model_(std::move(model)),
reader_(std::move(reader)),
model_str_(std::move(model_str)) {}
CalibrationWrapper::~CalibrationWrapper() = default;
PyObject* CalibrationWrapper::Prepare() {
TFLITE_PY_ENSURE_VALID_INTERPRETER();
TFLITE_PY_CHECK(interpreter_->AllocateTensors());
TFLITE_PY_CHECK(interpreter_->ResetVariableTensors());
Py_RETURN_NONE;
}
PyObject* CalibrationWrapper::Prepare(std::string signature_key) {
TFLITE_PY_ENSURE_VALID_INTERPRETER();
impl::SignatureRunner* runner =
interpreter_->GetSignatureRunner(signature_key.c_str());
if (runner == nullptr) {
PyErr_Format(PyExc_ValueError, "Invalid signature key: %s",
signature_key.c_str());
return nullptr;
}
TFLITE_PY_CHECK(runner->AllocateTensors());
TFLITE_PY_CHECK(interpreter_->ResetVariableTensors());
Py_RETURN_NONE;
}
PyObject* CalibrationWrapper::Prepare(PyObject* input_shapes,
std::string signature_key) {
TFLITE_PY_ENSURE_VALID_INTERPRETER();
if (!PyList_Check(input_shapes)) {
PyErr_Format(PyExc_ValueError,
"Invalid input shapes: expected shapes to be a list.");
return nullptr;
}
const int subgraph_index =
interpreter_->GetSubgraphIndexFromSignature(signature_key.c_str());
if (subgraph_index == -1) {
PyErr_Format(PyExc_ValueError, "Invalid signature key: %s",
signature_key.c_str());
return nullptr;
}
auto* subgraph = interpreter_->subgraph(subgraph_index);
const size_t inputs_size = PyList_Size(input_shapes);
if (inputs_size != subgraph->inputs().size()) {
PyErr_Format(PyExc_ValueError,
"Invalid input shapes: expected %ld items got %ld items.",
subgraph->inputs().size(), inputs_size);
return nullptr;
}
for (size_t i = 0; i < inputs_size; ++i) {
std::optional<std::vector<int>> dims =
ConvertInputShapeToVector(input_shapes, i);
if (!dims.has_value()) {
return nullptr;
}
int input_tensor_idx = subgraph->inputs()[i];
if (subgraph->ResizeInputTensor(input_tensor_idx, *dims) != kTfLiteOk) {
PyErr_Format(PyExc_ValueError, "Failed to resize %ld input tensor.", i);
return nullptr;
}
}
return Prepare(signature_key);
}
PyObject* CalibrationWrapper::Prepare(PyObject* input_shapes) {
TFLITE_PY_ENSURE_VALID_INTERPRETER();
if (!PyList_Check(input_shapes)) {
PyErr_Format(PyExc_ValueError,
"Invalid input shapes: expected shapes to be a list.");
return nullptr;
}
const size_t inputs_size = PyList_Size(input_shapes);
if (inputs_size != interpreter_->inputs().size()) {
PyErr_Format(PyExc_ValueError,
"Invalid input shapes: expected %ld items got %ld items.",
interpreter_->inputs().size(), inputs_size);
return nullptr;
}
for (size_t i = 0; i < inputs_size; ++i) {
std::optional<std::vector<int>> dims =
ConvertInputShapeToVector(input_shapes, i);
if (!dims.has_value()) {
return nullptr;
}
int input_tensor_idx = interpreter_->inputs()[i];
if (interpreter_->ResizeInputTensor(input_tensor_idx, *dims) != kTfLiteOk) {
PyErr_Format(PyExc_ValueError, "Failed to resize %ld input tensor.", i);
return nullptr;
}
}
return Prepare();
}
PyObject* CalibrationWrapper::FeedTensor(PyObject* input_value,
std::string signature_key) {
TFLITE_PY_ENSURE_VALID_INTERPRETER();
if (!PyList_Check(input_value)) {
PyErr_Format(PyExc_ValueError,
"Invalid input type: expected input to be a list.");
return nullptr;
}
const int subgraph_index =
interpreter_->GetSubgraphIndexFromSignature(signature_key.c_str());
if (subgraph_index == -1) {
PyErr_Format(PyExc_ValueError, "Invalid signature key: %s",
signature_key.c_str());
return nullptr;
}
const size_t inputs_size = PyList_Size(input_value);
auto* subgraph = interpreter_->subgraph(subgraph_index);
if (inputs_size != subgraph->inputs().size()) {
PyErr_Format(PyExc_ValueError,
"Invalid input size: expected %ld items got %ld items.",
subgraph->inputs().size(), inputs_size);
return nullptr;
}
for (size_t i = 0; i < inputs_size; ++i) {
PyObject* input = PyList_GetItem(input_value, i);
if (!input) {
return nullptr;
}
int input_tensor_idx = subgraph->inputs()[i];
if (!SetTensor(input_tensor_idx, input, signature_key)) {
return nullptr;
}
}
TFLITE_PY_CHECK(subgraph->Invoke());
Py_RETURN_NONE;
}
PyObject* CalibrationWrapper::FeedTensor(PyObject* input_value) {
TFLITE_PY_ENSURE_VALID_INTERPRETER();
if (!PyList_Check(input_value)) {
PyErr_Format(PyExc_ValueError,
"Invalid input type: expected input to be a list.");
return nullptr;
}
const size_t inputs_size = PyList_Size(input_value);
if (inputs_size != interpreter_->inputs().size()) {
PyErr_Format(PyExc_ValueError,
"Invalid input size: expected %ld items got %ld items.",
interpreter_->inputs().size(), inputs_size);
return nullptr;
}
for (size_t i = 0; i < inputs_size; ++i) {
PyObject* input = PyList_GetItem(input_value, i);
if (!input) {
return nullptr;
}
int input_tensor_idx = interpreter_->inputs()[i];
if (!SetTensor(input_tensor_idx, input)) {
return nullptr;
}
}
TFLITE_PY_CHECK(interpreter_->Invoke());
Py_RETURN_NONE;
}
PyObject* CalibrationWrapper::SetTensor(int index, PyObject* value,
std::string signature_key) {
TFLITE_PY_ENSURE_VALID_INTERPRETER();
std::unique_ptr<PyObject, PyDecrefDeleter> array_safe(
PyArray_FromAny(value, nullptr, 0, 0, NPY_ARRAY_CARRAY, nullptr));
if (!array_safe) {
PyErr_SetString(PyExc_ValueError,
"Failed to convert value into readable tensor.");
return nullptr;
}
PyArrayObject* array = reinterpret_cast<PyArrayObject*>(array_safe.get());
const int subgraph_index =
interpreter_->GetSubgraphIndexFromSignature(signature_key.c_str());
if (subgraph_index == -1) {
PyErr_Format(PyExc_ValueError, "Invalid signature key: %s",
signature_key.c_str());
return nullptr;
}
auto* subgraph = interpreter_->subgraph(subgraph_index);
const TfLiteTensor* tensor = subgraph->tensor(index);
if (TfLiteTypeFromPyArray(array) != tensor->type) {
PyErr_Format(PyExc_ValueError,
"Cannot set tensor: "
"Got value of type %s "
"but expected type %s for input %d, name: %s ",
TfLiteTypeGetName(TfLiteTypeFromPyArray(array)),
TfLiteTypeGetName(tensor->type), index, tensor->name);
return nullptr;
}
if (PyArray_NDIM(array) != tensor->dims->size) {
PyErr_Format(PyExc_ValueError,
"Cannot set tensor: Dimension count mismatch, expected %d "
"but found %d",
tensor->dims->size, PyArray_NDIM(array));
return nullptr;
}
std::vector<int> dims(PyArray_NDIM(array));
bool has_unknown_dims = false;
for (int j = 0; j < PyArray_NDIM(array); ++j) {
// Ensure the calibration data input shape is the same as the model input
// shape unless the dimension is unknown.
if (tensor->dims_signature != nullptr &&
tensor->dims_signature->size == tensor->dims->size &&
tensor->dims_signature->data[j] == -1) {
has_unknown_dims = true;
} else if (tensor->dims->data[j] != PyArray_SHAPE(array)[j]) {
PyErr_Format(PyExc_ValueError,
"Cannot set tensor: Size mismatch, expected %d for dim "
"%d but found %ld",
tensor->dims->data[j], j, PyArray_SHAPE(array)[j]);
return nullptr;
}
dims[j] = PyArray_SHAPE(array)[j];
}
// Resize the input tensor if there are unknown dimensions.
if (has_unknown_dims) {
// Does strict checking on the `ResizeInputTensor` call.
TFLITE_PY_CHECK(subgraph->ResizeInputTensorStrict(index, dims));
TFLITE_PY_CHECK(subgraph->AllocateTensors());
}
// Re-read the updated tensor after the allocation is done.
tensor = subgraph->tensor(index);
size_t size = PyArray_NBYTES(array);
if (tensor->type == kTfLiteString) {
DynamicBuffer buffer;
buffer.AddString(reinterpret_cast<const char*>(PyArray_BYTES(array)), size);
buffer.WriteToTensor(subgraph->tensor(index), /*new_shape=*/nullptr);
Py_RETURN_NONE;
}
if (size != tensor->bytes) {
PyErr_Format(PyExc_ValueError,
"numpy array had %zu bytes but expected %zu bytes.", size,
tensor->bytes);
return nullptr;
}
memcpy(tensor->data.raw, PyArray_DATA(array), size);
Py_RETURN_NONE;
}
PyObject* CalibrationWrapper::SetTensor(int index, PyObject* value) {
TFLITE_PY_ENSURE_VALID_INTERPRETER();
std::unique_ptr<PyObject, PyDecrefDeleter> array_safe(
PyArray_FromAny(value, nullptr, 0, 0, NPY_ARRAY_CARRAY, nullptr));
if (!array_safe) {
PyErr_SetString(PyExc_ValueError,
"Failed to convert value into readable tensor.");
return nullptr;
}
PyArrayObject* array = reinterpret_cast<PyArrayObject*>(array_safe.get());
const TfLiteTensor* tensor = interpreter_->tensor(index);
if (TfLiteTypeFromPyArray(array) != tensor->type) {
PyErr_Format(PyExc_ValueError,
"Cannot set tensor: "
"Got value of type %s "
"but expected type %s for input %d, name: %s ",
TfLiteTypeGetName(TfLiteTypeFromPyArray(array)),
TfLiteTypeGetName(tensor->type), index, tensor->name);
return nullptr;
}
if (PyArray_NDIM(array) != tensor->dims->size) {
PyErr_Format(
PyExc_ValueError,
"Cannot set tensor: Dimension count mismatch, expected %d but found %d",
tensor->dims->size, PyArray_NDIM(array));
return nullptr;
}
std::vector<int> dims(PyArray_NDIM(array));
bool has_unknown_dims = false;
for (int j = 0; j < PyArray_NDIM(array); ++j) {
// Ensure the calibration data input shape is the same as the model input
// shape unless the dimension is unknown.
if (tensor->dims_signature != nullptr &&
tensor->dims_signature->size == tensor->dims->size &&
tensor->dims_signature->data[j] == -1) {
has_unknown_dims = true;
} else if (tensor->dims->data[j] != PyArray_SHAPE(array)[j]) {
PyErr_Format(PyExc_ValueError,
"Cannot set tensor: Size mismatch, expected %d for dim "
"%d but found %ld",
tensor->dims->data[j], j, PyArray_SHAPE(array)[j]);
return nullptr;
}
dims[j] = PyArray_SHAPE(array)[j];
}
// Resize the input tensor if there are unknown dimensions.
if (has_unknown_dims) {
// Does strict checking on the `ResizeInputTensor` call.
TFLITE_PY_CHECK(interpreter_->ResizeInputTensorStrict(index, dims));
TFLITE_PY_CHECK(interpreter_->AllocateTensors());
}
// Re-read the updated tensor after the allocation is done.
tensor = interpreter_->tensor(index);
size_t size = PyArray_NBYTES(array);
if (tensor->type == kTfLiteString) {
DynamicBuffer buffer;
buffer.AddString(reinterpret_cast<const char*>(PyArray_BYTES(array)), size);
buffer.WriteToTensor(interpreter_->tensor(index), /*new_shape=*/nullptr);
Py_RETURN_NONE;
}
if (size != tensor->bytes) {
PyErr_Format(PyExc_ValueError,
"numpy array had %zu bytes but expected %zu bytes.", size,
tensor->bytes);
return nullptr;
}
memcpy(tensor->data.raw, PyArray_DATA(array), size);
Py_RETURN_NONE;
}
PyObject* CalibrationWrapper::Calibrate() {
const auto [base_buffer, offset_buffer] =
SplitOffsetBuffer(/*model_buffer=*/absl::string_view(
reinterpret_cast<const char*>(model_->allocation()->base()),
model_->allocation()->bytes()));
auto tflite_model = CreateMutableModel(*model_->GetModel());
reader_->AddCalibrationToModel(tflite_model.get(), /*update=*/false);
flatbuffers::FlatBufferBuilder builder;
auto loc = Model::Pack(builder, tflite_model.get());
FinishModelBuffer(builder, loc);
const int64_t result_base_buffer_size = builder.GetSize();
const int64_t offset_diff =
result_base_buffer_size - static_cast<int64_t>(base_buffer.size());
const std::string updated_result_base_buffer = UpdateBufferOffsets(
/*base_buffer=*/absl::string_view(
reinterpret_cast<const char*>(builder.GetCurrentBufferPointer()),
result_base_buffer_size),
offset_diff);
const std::string result_buffer =
MergeOffsetBuffer(updated_result_base_buffer, offset_buffer);
return ConvertToPyString(result_buffer.data(), result_buffer.size());
}
PyObject* CalibrationWrapper::QuantizeModel(int input_py_type,
int output_py_type,
bool allow_float,
int activations_py_type,
int bias_py_type) {
return QuantizeModel(
input_py_type, output_py_type, allow_float, activations_py_type,
bias_py_type,
/*disable_per_channel=*/false,
/*disable_per_channel_quantization_for_dense_layers=*/false);
}
PyObject* CalibrationWrapper::QuantizeModel(
int input_py_type, int output_py_type, bool allow_float,
int activations_py_type, int bias_py_type, bool disable_per_channel,
bool disable_per_channel_quantization_for_dense_layers) {
if (NoOpModel(*model_)) {
return ConvertToPyString(model_str_->data(), model_str_->size());
}
TfLiteType input_type = TfLiteTypeFromPyType(input_py_type);
TfLiteType output_type = TfLiteTypeFromPyType(output_py_type);
TfLiteType activations_type = TfLiteTypeFromPyType(activations_py_type);
TfLiteType bias_type = TfLiteTypeFromPyType(bias_py_type);
if (input_type == kTfLiteNoType || output_type == kTfLiteNoType) {
PyErr_SetString(PyExc_ValueError,
"Input/output type cannot be kTfLiteNoType");
return nullptr;
}
auto tflite_model = CreateMutableModel(*model_->GetModel());
reader_->AddCalibrationToModel(tflite_model.get(), /*update=*/false);
flatbuffers::FlatBufferBuilder builder;
auto status = kTfLiteOk;
status = QuantizeModelAllOperators(
&builder, tflite_model.get(), TfLiteTypeToSchemaType(input_type),
TfLiteTypeToSchemaType(output_type), allow_float,
TfLiteTypeToSchemaType(activations_type),
TfLiteTypeToSchemaType(bias_type), disable_per_channel,
disable_per_channel_quantization_for_dense_layers, error_reporter_.get());
if (status != kTfLiteOk) {
error_reporter_->exception();
return nullptr;
}
return ConvertToPyString(
reinterpret_cast<const char*>(builder.GetCurrentBufferPointer()),
builder.GetSize());
}
PyObject* CalibrationWrapper::QuantizeModel(int input_py_type,
int output_py_type,
bool allow_float,
const char* operator_output_name) {
string op_name = std::string(operator_output_name);
TfLiteType input_type = TfLiteTypeFromPyType(input_py_type);
TfLiteType output_type = TfLiteTypeFromPyType(output_py_type);
if (input_type == kTfLiteNoType || output_type == kTfLiteNoType) {
PyErr_SetString(PyExc_ValueError,
"Input/output type cannot be kTfLiteNoType");
return nullptr;
}
auto tflite_model = CreateMutableModel(*model_->GetModel());
reader_->AddCalibrationToModel(tflite_model.get(), /*update=*/false);
flatbuffers::FlatBufferBuilder builder;
auto status = optimize::QuantizeModel(
&builder, tflite_model.get(), TfLiteTypeToSchemaType(input_type),
TfLiteTypeToSchemaType(output_type), allow_float, {op_name},
/*activations_type=*/TensorType_INT8, /*bias_type=*/TensorType_INT32,
error_reporter_.get());
if (status != kTfLiteOk) {
error_reporter_->exception();
return nullptr;
}
return ConvertToPyString(
reinterpret_cast<const char*>(builder.GetCurrentBufferPointer()),
builder.GetSize());
}
/*static*/ CalibrationWrapper* CalibrationWrapper::CreateWrapperCPPFromBuffer(
PyObject* data, const std::vector<std::string>& registerers_by_name,
const std::vector<std::function<void(uintptr_t)>>& registerers_by_func,
std::string* error_msg) {
char* buf = nullptr;
Py_ssize_t length;
std::unique_ptr<PythonErrorReporter> error_reporter(new PythonErrorReporter);
ImportNumpy();
if (ConvertFromPyString(data, &buf, &length) == -1) {
*error_msg = "Failed to convert from python string";
return nullptr;
}
std::unique_ptr<FlatBufferModel> model =
FlatBufferModel::BuildFromBuffer(buf, length, error_reporter.get());
if (!model) {
*error_msg = "Invalid model";
return nullptr;
}
auto resolver = std::make_unique<BuiltinOpResolver>();
for (const auto& registerer : registerers_by_name) {
if (!RegisterCustomOpByName(registerer.c_str(), resolver.get())) {
*error_msg =
absl::StrFormat("Looking up symbol '%s' failed with error '%s'.",
registerer.c_str(), SharedLibrary::GetError());
return nullptr;
}
}
for (const auto& registerer : registerers_by_func) {
registerer(reinterpret_cast<uintptr_t>(resolver.get()));
}
std::unique_ptr<Interpreter> interpreter;
std::unique_ptr<CalibrationReader> reader;
auto status =
BuildLoggingInterpreter(*model, *resolver, &interpreter, &reader);
if (status != kTfLiteOk) {
*error_msg = error_reporter->message();
return nullptr;
}
auto model_str = std::make_unique<std::string>(buf, length);
// If we are not going to use this string during quantization, reset the
// pointer and release the memory.
if (!NoOpModel(*model)) {
model_str.reset();
}
auto wrapper = new CalibrationWrapper(
std::move(interpreter), std::move(resolver), std::move(error_reporter),
std::move(model), std::move(reader), std::move(model_str));
return wrapper;
}
} // namespace calibration_wrapper
} // namespace tflite
tensorflow/lite/python/optimize/calibration_wrapper.h
+138
View File
@@ -0,0 +1,138 @@
/* 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_PYTHON_OPTIMIZE_CALIBRATION_WRAPPER_H_
#define TENSORFLOW_LITE_PYTHON_OPTIMIZE_CALIBRATION_WRAPPER_H_
#include <functional>
#include <memory>
#include <string>
#include <vector>
// Place `<locale>` before <Python.h> to avoid build failures in macOS.
#include <locale>
// The empty line above is on purpose as otherwise clang-format will
// automatically move <Python.h> before <locale>.
#include <Python.h>
#include "tensorflow/lite/core/interpreter.h"
// We forward declare TFLite classes here to avoid exposing them to SWIG.
namespace tflite {
namespace ops {
namespace builtin {
class BuiltinOpResolver;
} // namespace builtin
} // namespace ops
namespace impl {
class FlatBufferModel;
}
namespace interpreter_wrapper {
class PythonErrorReporter;
} // namespace interpreter_wrapper
namespace optimize {
namespace calibration {
class CalibrationReader;
} // namespace calibration
} // namespace optimize
namespace calibration_wrapper {
PyObject* AddIntermediateTensors(PyObject* data);
class CalibrationWrapper {
public:
// SWIG caller takes ownership of pointer.
static CalibrationWrapper* CreateWrapperCPPFromBuffer(
PyObject* data, const std::vector<std::string>& registerers_by_name,
const std::vector<std::function<void(uintptr_t)>>& registerers_by_func,
std::string* error_msg);
~CalibrationWrapper();
// Allocates the primary subgraph's tensors.
PyObject* Prepare();
// Allocates the tensors of the given signature, defined by the signature
// key.
PyObject* Prepare(std::string signature_key);
// Allocates the primary subgraph's tensors with the given input shapes.
PyObject* Prepare(PyObject* input_shapes);
// Allocates the tensors of the given signature with the given input
// shapes, defined by the signature key.
PyObject* Prepare(PyObject* input_shapes, std::string signature_key);
// Sets the given input tensors to the primary subgraph.
PyObject* FeedTensor(PyObject* input_value);
// Sets the given input tensor to the given signature, defined by the
// signature key.
PyObject* FeedTensor(PyObject* input_value, std::string signature_key);
// Allows quantizing only the operator that produces the tensor.
PyObject* QuantizeModel(int input_py_type, int output_py_type,
bool allow_float, int activations_py_type,
int bias_py_type);
// Allows quantizing only the operator that produces the tensor with name
// operator_output_name. (This can be used to help debug.).
// TODO(suharshs): Allow providing multiple names.
PyObject* QuantizeModel(int input_py_type, int output_py_type,
bool allow_float, const char* operator_output_name);
// Disables per-channel quantization, can be used to produce smaller
// models but may cause accuracy issues.
PyObject* QuantizeModel(
int input_py_type, int output_py_type, bool allow_float,
int activations_py_type, int bias_py_type, bool disable_per_channel,
bool disable_per_channel_quantization_for_dense_layers);
// Writes the in-memory calibration results to the model flatbuffer. The
// produced model is as same as the original input model, but the min/max
// in the quantization field.
PyObject* Calibrate();
private:
// CalibrationWrapper is not copyable or assignable. We avoid the use of
// CalibrationWrapper() = delete here for SWIG compatibility.
CalibrationWrapper(
std::unique_ptr<Interpreter> interpreter,
std::unique_ptr<ops::builtin::BuiltinOpResolver> resolver,
std::unique_ptr<interpreter_wrapper::PythonErrorReporter> error_reporter,
std::unique_ptr<impl::FlatBufferModel> model,
std::unique_ptr<optimize::calibration::CalibrationReader> reader,
std::unique_ptr<std::string> model_str_);
CalibrationWrapper(const CalibrationWrapper& rhs);
PyObject* SetTensor(int index, PyObject* value);
PyObject* SetTensor(int index, PyObject* value, std::string signature_key);
std::unique_ptr<Interpreter> interpreter_;
std::unique_ptr<interpreter_wrapper::PythonErrorReporter> error_reporter_;
std::unique_ptr<ops::builtin::BuiltinOpResolver> resolver_;
std::unique_ptr<impl::FlatBufferModel> model_;
std::unique_ptr<optimize::calibration::CalibrationReader> reader_;
std::unique_ptr<std::string> model_str_;
};
} // namespace calibration_wrapper
} // namespace tflite
#endif // TENSORFLOW_LITE_PYTHON_OPTIMIZE_CALIBRATION_WRAPPER_H_
tensorflow/lite/python/optimize/calibration_wrapper_pybind11.cc
+102
View File
@@ -0,0 +1,102 @@
/* Copyright 2020 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.
==============================================================================*/
#include <cstdint>
#include <functional>
#include <stdexcept>
#include <string>
#include <vector>
#include "pybind11/functional.h" // from @pybind11
#include "pybind11/pybind11.h" // from @pybind11
#include "pybind11/pytypes.h" // from @pybind11
#include "pybind11/stl.h" // from @pybind11
#include "tensorflow/lite/python/optimize/calibration_wrapper.h"
#include "tensorflow/python/lib/core/pybind11_lib.h"
namespace py = pybind11;
using tflite::calibration_wrapper::AddIntermediateTensors;
using tflite::calibration_wrapper::CalibrationWrapper;
PYBIND11_MODULE(_pywrap_tensorflow_lite_calibration_wrapper, m) {
m.doc() = R"pbdoc(
_pywrap_tensorflow_lite_calibration_wrapper
-----
)pbdoc";
m.def("AddIntermediateTensors", [](py::handle& data) {
return tensorflow::PyoOrThrow(AddIntermediateTensors(data.ptr()));
});
py::class_<CalibrationWrapper>(m, "CalibrationWrapper")
.def(py::init([](py::handle& data,
const std::vector<std::string>& registerers_by_name,
const std::vector<std::function<void(uintptr_t)>>&
registerers_by_func) {
std::string error;
auto* wrapper = ::CalibrationWrapper::CreateWrapperCPPFromBuffer(
data.ptr(), registerers_by_name, registerers_by_func, &error);
if (!wrapper) {
throw std::invalid_argument(error); // throws ValueError in Python
}
return wrapper;
}))
.def("Prepare",
[](CalibrationWrapper& self, py::handle& input_shapes,
std::string signature_key) {
return tensorflow::PyoOrThrow(
self.Prepare(input_shapes.ptr(), signature_key));
})
.def("Prepare",
[](CalibrationWrapper& self, py::handle& input_shapes) {
return tensorflow::PyoOrThrow(self.Prepare(input_shapes.ptr()));
})
.def("Prepare",
[](CalibrationWrapper& self, std::string signature_key) {
return tensorflow::PyoOrThrow(self.Prepare(signature_key));
})
.def("Prepare",
[](CalibrationWrapper& self) {
return tensorflow::PyoOrThrow(self.Prepare());
})
.def("FeedTensor",
[](CalibrationWrapper& self, py::handle& input_value,
std::string signature_key) {
return tensorflow::PyoOrThrow(
self.FeedTensor(input_value.ptr(), signature_key));
})
.def("FeedTensor",
[](CalibrationWrapper& self, py::handle& input_value) {
return tensorflow::PyoOrThrow(self.FeedTensor(input_value.ptr()));
})
.def("QuantizeModel",
[](CalibrationWrapper& self, int input_py_type, int output_py_type,
bool allow_float, int activations_py_type, int bias_py_type,
bool disable_per_channel,
bool disable_per_channel_quantization_for_dense_layers) {
return tensorflow::PyoOrThrow(self.QuantizeModel(
input_py_type, output_py_type, allow_float,
activations_py_type, bias_py_type, disable_per_channel,
disable_per_channel_quantization_for_dense_layers));
})
.def("QuantizeModel",
[](CalibrationWrapper& self, int input_py_type, int output_py_type,
bool allow_float, const char* operator_output_name) {
return tensorflow::PyoOrThrow(
self.QuantizeModel(input_py_type, output_py_type, allow_float,
operator_output_name));
})
.def("Calibrate", [](CalibrationWrapper& self) {
return tensorflow::PyoOrThrow(self.Calibrate());
});
}
tensorflow/lite/python/optimize/calibrator.py
+259
View File
@@ -0,0 +1,259 @@
# 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.
# ==============================================================================
"""Python wrapper for post training quantization with calibration."""
import numpy as np
from tensorflow.lite.python.convert_phase import Component
from tensorflow.lite.python.convert_phase import convert_phase
from tensorflow.lite.python.convert_phase import SubComponent
from tensorflow.lite.python.interpreter import Interpreter
from tensorflow.python.framework import dtypes
from tensorflow.python.util.lazy_loader import LazyLoader
# Lazy load since some of the performance benchmark skylark rules
# break dependencies. Must use double quotes to match code internal rewrite
# rule.
_calibration_wrapper = LazyLoader(
"_calibration_wrapper",
globals(),
(
"tensorflow.lite.python.optimize."
"_pywrap_tensorflow_lite_calibration_wrapper"
),
)
def add_intermediate_tensors(model_content):
"""Adds intermediate tensors to fused op if needed."""
return _calibration_wrapper.AddIntermediateTensors(model_content)
class Calibrator:
"""Calibrates a floating point model and then quantizes it.
This is an internal class, not a public interface.
"""
def __init__(
self,
model_content,
custom_op_registerers_by_name=None,
custom_op_registerers_by_func=None,
):
"""Constructor.
Args:
model_content: Content of a TF-Lite Flatbuffer file.
custom_op_registerers_by_name: List of str (symbol names) that take a
pointer to a MutableOpResolver and register custom ops.
custom_op_registerers_by_func: List of functions that take a pointer to a
MutableOpResolver and register custom ops.
Raises:
ValueError: If the calibrator was unable to open the model.
"""
if not model_content:
raise ValueError("`model_content` must be specified.")
if custom_op_registerers_by_name is None:
custom_op_registerers_by_name = []
if custom_op_registerers_by_func is None:
custom_op_registerers_by_func = []
try:
self._calibrator = _calibration_wrapper.CalibrationWrapper(
model_content,
custom_op_registerers_by_name,
custom_op_registerers_by_func,
)
self._model_content = model_content
except Exception as e:
raise ValueError("Failed to parse the model: %s." % e)
if not self._calibrator:
raise ValueError("Failed to parse the model.")
self._interpreter = None
def _create_input_array_from_dict(self, signature_key, inputs):
input_array = []
signature_runner = self._interpreter.get_signature_runner(signature_key)
input_details = sorted(
signature_runner.get_input_details().items(),
key=lambda item: item[1]["index"],
)
for input_name, _ in input_details:
input_array.append(inputs[input_name])
return input_array
def _feed_tensors(self, dataset_gen, resize_input):
"""Feed tensors to the calibrator."""
initialized = {}
for sample in dataset_gen():
if isinstance(sample, tuple):
if not isinstance(sample[1], dict):
raise ValueError(
"You need to provide either a dictionary with input "
"names and values in the second argument in the "
"tuple"
)
# Convert signature based inputs to the tensor index based data.
if self._interpreter is None:
self._interpreter = Interpreter(model_content=self._model_content)
signature_key = sample[0]
input_array = self._create_input_array_from_dict(
signature_key, sample[1]
)
elif isinstance(sample, dict):
# Convert signature based inputs to the tensor index based data.
if self._interpreter is None:
self._interpreter = Interpreter(model_content=self._model_content)
signature_key = None
input_array = self._create_input_array_from_dict(None, sample)
elif isinstance(sample, list):
signature_key = None
input_array = sample
else:
raise ValueError(
"You need to provide either a dictionary with input "
"names and values, a tuple with signature key and a "
"dictionary with input names and values, or an array "
"with input values in the order of input tensors of "
"the graph in the representative_dataset function. "
"Unsupported value from dataset: {}.".format(sample)
)
if signature_key not in initialized:
initialized[signature_key] = True
if resize_input:
if signature_key is not None:
self._calibrator.Prepare(
[list(s.shape) for s in input_array], signature_key
)
else:
self._calibrator.Prepare([list(s.shape) for s in input_array])
else:
if signature_key is not None:
self._calibrator.Prepare(signature_key)
else:
self._calibrator.Prepare()
if signature_key is not None:
self._calibrator.FeedTensor(input_array, signature_key)
else:
self._calibrator.FeedTensor(input_array)
@convert_phase(
Component.OPTIMIZE_TFLITE_MODEL,
SubComponent.QUANTIZE_USING_DEPRECATED_QUANTIZER,
)
def calibrate_and_quantize(
self,
dataset_gen,
input_type,
output_type,
allow_float,
activations_type=dtypes.int8,
bias_type=dtypes.int32,
resize_input=True,
disable_per_channel=False,
disable_per_channel_quantization_for_dense_layers=False,
):
"""Calibrates the model with specified generator and then quantizes it.
The input shapes of the calibrator are resized with the calibration data if
`resize_input` is set.
Returns:
A quantized model.
Args:
dataset_gen: A generator that generates calibration samples.
input_type: A tf.dtype representing the desired real-value input type.
output_type: A tf.dtype representing the desired real-value output type.
allow_float: A boolean. False if the resulting model cannot perform float
computation, useful when targeting an integer-only backend. If False, an
error will be thrown if an operation cannot be quantized, otherwise the
model will fallback to float ops.
activations_type: A tf.dtype representing the desired type for
activations.
bias_type: A tf.dtype representing the desired type for bias.
resize_input: A boolean. True if the shape of the sample data is different
from the input.
disable_per_channel: A boolean. True if disabling per-channel
quantization.
disable_per_channel_quantization_for_dense_layers: A boolean. True if
disabling per-channel quantization only in Dense layers.
"""
self._feed_tensors(dataset_gen, resize_input)
return self._calibrator.QuantizeModel(
np.dtype(input_type.as_numpy_dtype()).num,
np.dtype(output_type.as_numpy_dtype()).num,
allow_float,
np.dtype(activations_type.as_numpy_dtype()).num,
np.dtype(bias_type.as_numpy_dtype()).num,
disable_per_channel,
disable_per_channel_quantization_for_dense_layers,
)
@convert_phase(
Component.OPTIMIZE_TFLITE_MODEL,
SubComponent.QUANTIZE_USING_DEPRECATED_QUANTIZER,
)
def calibrate_and_quantize_single(
self,
dataset_gen,
input_type,
output_type,
allow_float,
op_output_name,
resize_input=True,
):
"""Calibrates the model with specified generator and then quantizes it.
Only the single op with output op_output_name will be quantized.
The input shapes of the calibrator are resized with the calibration data.
Returns:
A quantized model.
Args:
dataset_gen: A generator that generates calibration samples.
input_type: A tf.dtype representing the desired real-value input type.
output_type: A tf.dtype representing the desired real-value output type.
allow_float: A boolean. False if the resulting model cannot perform float
computation, useful when targeting an integer-only backend. If False, an
error will be thrown if an operation cannot be quantized, otherwise the
model will fallback to float ops.
op_output_name: A string, only this op will be quantized.
resize_input: A boolean. True if the shape of the sample data is different
from the input.
"""
self._feed_tensors(dataset_gen, resize_input)
return self._calibrator.QuantizeModel(
np.dtype(input_type.as_numpy_dtype()).num,
np.dtype(output_type.as_numpy_dtype()).num,
allow_float,
op_output_name,
)
@convert_phase(Component.OPTIMIZE_TFLITE_MODEL, SubComponent.CALIBRATE)
def calibrate(self, dataset_gen):
"""Calibrates the model with specified generator.
Returns:
A model with min and max calibration stats.
Args:
dataset_gen: A generator that generates calibration samples.
"""
self._feed_tensors(dataset_gen, resize_input=True)
return self._calibrator.Calibrate()
tensorflow/lite/python/optimize/calibrator_test.py
+311
View File
@@ -0,0 +1,311 @@
# 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.
# ==============================================================================
"""Tests for Calibrator."""
from absl.testing import parameterized
import numpy as np
import tensorflow as tf
from tensorflow.lite.python import lite
from tensorflow.lite.python import schema_py_generated as schema_fb
from tensorflow.lite.python.optimize import calibrator as _calibrator
from tensorflow.lite.tools import flatbuffer_utils
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import test_util
from tensorflow.python.platform import resource_loader
from tensorflow.python.platform import test
def _uses_buffer_offset(model: schema_fb.ModelT) -> bool:
"""Determines whether the model is using an offset buffer.
Args:
model: A TFLite model.
Returns:
True iff the model is using offset buffers. Offset buffers are enabled by
the flag `_experimental_use_buffer_offset`.
"""
if not model.metadata:
return False
return any(
map(
lambda metadata: metadata.name.decode('utf-8') == 'buffer_location',
model.metadata,
)
)
class CalibratorTest(test_util.TensorFlowTestCase, parameterized.TestCase):
@parameterized.named_parameters(
# Activation type Int8
('UseActivationTypeInt8', dtypes.int8),
# Activation type Int16
('UseActivationTypeInt16', dtypes.int16),
)
def test_calibration_with_quantization(self, activations_type):
model_path = resource_loader.get_path_to_datafile(
'test_data/mobilenet_like_model.bin'
)
float_model = open(model_path, 'rb').read()
quantizer = _calibrator.Calibrator(float_model)
# Input generator for the model.
def input_gen():
for _ in range(10):
yield [np.ones(shape=(1, 5, 5, 3), dtype=np.float32)]
quantized_model = quantizer.calibrate_and_quantize(
input_gen, dtypes.float32, dtypes.float32, False, activations_type
)
self.assertIsNotNone(quantized_model)
@parameterized.named_parameters(
# Activation type Int8
('UseActivationTypeInt8', dtypes.int8),
# Activation type Int16
('UseActivationTypeInt16', dtypes.int16),
)
def test_calibration_with_quantization_allow_float(self, activations_type):
model_path = resource_loader.get_path_to_datafile(
'test_data/mobilenet_like_model.bin'
)
float_model = open(model_path, 'rb').read()
quantizer = _calibrator.Calibrator(float_model)
# Input generator for the model.
def input_gen():
for _ in range(10):
yield [np.ones(shape=(1, 5, 5, 3), dtype=np.float32)]
quantized_model = quantizer.calibrate_and_quantize(
input_gen, dtypes.float32, dtypes.float32, True, activations_type
)
self.assertIsNotNone(quantized_model)
def test_calibration_with_quantization_single_op(self):
model_path = resource_loader.get_path_to_datafile(
'test_data/mobilenet_like_model.bin'
)
float_model = open(model_path, 'rb').read()
quantizer = _calibrator.Calibrator(float_model)
# Input generator for the model.
def input_gen():
for _ in range(10):
yield [np.ones(shape=(1, 5, 5, 3), dtype=np.float32)]
quantized_model = quantizer.calibrate_and_quantize_single(
input_gen, dtypes.float32, dtypes.float32, True, 'conv2d_8/BiasAdd'
)
self.assertIsNotNone(quantized_model)
def test_calibration_with_string_input(self):
model_path = resource_loader.get_path_to_datafile(
'test_data/string_input_flex_model.bin'
)
with open(model_path, 'rb') as fp:
model_with_string_input = fp.read()
quantizer = _calibrator.Calibrator(model_with_string_input)
# Input generator for the model.
def input_gen():
for i in range(10):
yield [np.array('Test' + str(i))]
quantized_model = quantizer.calibrate_and_quantize_single(
input_gen, dtypes.float32, dtypes.float32, True, 'Identity'
)
self.assertIsNotNone(quantized_model)
@parameterized.named_parameters(
# Activation type Int8
('UseActivationTypeInt8 - EnableMlirQuantizer', dtypes.int8),
# Activation type Int16
('UseActivationTypeInt16 - DisableEnableMlirQuantizer', dtypes.int16),
)
def test_calibration_with_quantization_multiple_inputs(
self, activations_type
):
# Load multi add model from test data.
# This model has 4 inputs of size (1, 8, 8, 3).
model_path = resource_loader.get_path_to_datafile(
'../../testdata/multi_add.bin'
)
float_model = open(model_path, 'rb').read()
quantizer = _calibrator.Calibrator(float_model)
# Input generator for the model.
def input_gen():
for _ in range(10):
yield [np.ones(shape=(1, 8, 8, 3), dtype=np.float32) for _ in range(4)]
quantized_model = quantizer.calibrate_and_quantize(
input_gen, dtypes.float32, dtypes.float32, False, activations_type
)
self.assertIsNotNone(quantized_model)
def test_invalid_model_buffer(self):
float_model = b'\0' * 100
with self.assertRaisesRegex(ValueError, 'Failed to parse the model'):
_calibrator.Calibrator(float_model)
# TODO(fengliuai): enable mlir quantizer
def test_empty_calibrator_gen(self):
model_path = resource_loader.get_path_to_datafile(
'test_data/mobilenet_like_model.bin'
)
float_model = open(model_path, 'rb').read()
quantizer = _calibrator.Calibrator(float_model)
def empty_input_gen():
for i in ():
yield i
with self.assertRaises(RuntimeError):
quantizer.calibrate_and_quantize(
empty_input_gen, dtypes.float32, dtypes.float32, False
)
def test_invalid_shape_calibrator_gen(self):
model_path = resource_loader.get_path_to_datafile(
'test_data/mobilenet_like_model.bin'
)
float_model = open(model_path, 'rb').read()
quantizer = _calibrator.Calibrator(float_model)
# Input generator with incorrect shape.
def input_gen():
for _ in range(10):
yield [np.ones(shape=(1, 2, 2, 3), dtype=np.float32)]
with self.assertRaisesRegex(ValueError, 'Size mismatch'):
quantizer.calibrate_and_quantize(
input_gen,
dtypes.float32,
dtypes.float32,
False,
activations_type=dtypes.int8,
bias_type=dtypes.int32,
resize_input=False,
)
def test_invalid_type_calibrator_gen(self):
model_path = resource_loader.get_path_to_datafile(
'test_data/mobilenet_like_model.bin'
)
float_model = open(model_path, 'rb').read()
quantizer = _calibrator.Calibrator(float_model)
# Input generator with incorrect type.
def input_gen():
for _ in range(10):
yield [np.ones(shape=(1, 5, 5, 3), dtype=np.int32)]
with self.assertRaises(ValueError):
quantizer.calibrate_and_quantize(
input_gen, dtypes.float32, dtypes.float32, False, dtypes.int8
)
def test_calibration(self):
model_path = resource_loader.get_path_to_datafile(
'test_data/mobilenet_like_model.bin'
)
float_model = open(model_path, 'rb').read()
quantizer = _calibrator.Calibrator(float_model)
# Input generator for the model.
def input_gen():
for _ in range(10):
yield [np.ones(shape=(1, 5, 5, 3), dtype=np.float32)]
quantized_model = quantizer.calibrate(input_gen)
self.assertIsNotNone(quantized_model)
def test_add_intermediate_tensors(self):
model_path = resource_loader.get_path_to_datafile(
'test_data/mobilenet_like_model.bin'
)
model = open(model_path, 'rb').read()
added_model = _calibrator.add_intermediate_tensors(model)
self.assertIsNotNone(added_model)
def test_calibrate_model_with_offset_buffer(self):
# Define a simple model to run calibration with.
class MatMulModel(tf.Module):
def __init__(self):
# Use ones for predictable calibration results.
self.filter = np.ones((4, 3)).astype(np.float32)
@tf.function(
input_signature=[tf.TensorSpec(shape=(1, 4), dtype=dtypes.float32)]
)
def __call__(self, input_tensor: tf.Tensor) -> tf.Tensor:
output_tensor = tf.linalg.matmul(input_tensor, self.filter)
return {'output': output_tensor}
model = MatMulModel()
saved_model_path = self.create_tempdir().full_path
tf.saved_model.save(model, saved_model_path)
converter = lite.TFLiteConverter.from_saved_model(saved_model_path)
# Enable the use of buffer offsets.
# pylint: disable=protected-access
converter._experimental_use_buffer_offset = True
# pylint: enable=protected-access
converter.exclude_conversion_metadata = True
model_serialized = converter.convert()
model = flatbuffer_utils.convert_bytearray_to_object(model_serialized)
self.assertTrue(_uses_buffer_offset(model))
quantizer = _calibrator.Calibrator(model_serialized)
# Input generator for the model.
def input_gen():
for _ in range(2):
yield [np.array([1.0, 1.0, 1.0, 1.0], dtype=np.float32)]
calibrated_model_serialized = quantizer.calibrate(input_gen)
self.assertIsNotNone(calibrated_model_serialized)
calibrated_model = flatbuffer_utils.convert_bytearray_to_object(
calibrated_model_serialized
)
self.assertTrue(_uses_buffer_offset(calibrated_model))
# Confirm that the tensors are correctly calibrated.
subgraph = calibrated_model.subgraphs[0]
matmul_input_tensor = subgraph.tensors[0]
self.assertAllClose(matmul_input_tensor.quantization.min, [1.0])
self.assertAllClose(matmul_input_tensor.quantization.max, [1.0])
matmul_filter_tensor = subgraph.tensors[1]
self.assertAllClose(matmul_filter_tensor.quantization.min, [1.0])
self.assertAllClose(matmul_filter_tensor.quantization.max, [1.0])
# The matmul is performed with all ones so the output is expected to be 4s.
matmul_output_tensor = subgraph.tensors[2]
self.assertAllClose(matmul_output_tensor.quantization.min, [4.0])
self.assertAllClose(matmul_output_tensor.quantization.max, [4.0])
if __name__ == '__main__':
test.main()
tensorflow/lite/python/optimize/test_data/mobilenet_like_model.bin
BIN
View File
Binary file not shown.
tensorflow/lite/python/optimize/test_data/string_input_flex_model.bin
BIN
View File
Binary file not shown.
tensorflow/lite/python/pywrap_tflite_common.json
+8
View File
@@ -0,0 +1,8 @@
{
"global": [
"Wrapped_PyInit_*"
],
"local": [
"*"
]
}
tensorflow/lite/python/pywrap_tflite_common.lds
+7
View File
@@ -0,0 +1,7 @@
{
global:
Wrapped_PyInit_*;
local:
*;
};
tensorflow/lite/python/pywrap_tflite_common_darwin.lds
+1
View File
@@ -0,0 +1 @@
*Wrapped_PyInit_*
tensorflow/lite/python/schema_util.py
+45
View File
@@ -0,0 +1,45 @@
# Copyright 2020 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.
# ==============================================================================
"""Schema utilities to get builtin code from operator code."""
from tensorflow.python.util import all_util
def get_builtin_code_from_operator_code(opcode):
"""Return the builtin code of the given operator code.
The following method is introduced to resolve op builtin code shortage
problem. The new builtin operator will be assigned to the extended builtin
code field in the flatbuffer schema. Those methods helps to hide builtin code
details.
Args:
opcode: Operator code.
Returns:
The builtin code of the given operator code.
"""
# Access BuiltinCode() method first if available.
if hasattr(opcode, 'BuiltinCode') and callable(opcode.BuiltinCode):
return max(opcode.BuiltinCode(), opcode.DeprecatedBuiltinCode())
return max(opcode.builtinCode, opcode.deprecatedBuiltinCode)
_allowed_symbols = [
'get_builtin_code_from_operator_code',
]
all_util.remove_undocumented(__name__, _allowed_symbols)
tensorflow/lite/python/test_util.py
+53
View File
@@ -0,0 +1,53 @@
# Copyright 2020 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.
# ==============================================================================
"""Functions used by multiple tflite test files."""
from tensorflow.lite.python import schema_py_generated as schema_fb
from tensorflow.lite.python import schema_util
from tensorflow.lite.tools import visualize
def get_ops_list(model_data):
"""Returns a set of ops in the tflite model data."""
model = schema_fb.Model.GetRootAsModel(model_data, 0)
op_set = set()
for subgraph_idx in range(model.SubgraphsLength()):
subgraph = model.Subgraphs(subgraph_idx)
for op_idx in range(subgraph.OperatorsLength()):
op = subgraph.Operators(op_idx)
opcode = model.OperatorCodes(op.OpcodeIndex())
builtin_code = schema_util.get_builtin_code_from_operator_code(opcode)
if builtin_code == schema_fb.BuiltinOperator.CUSTOM:
opname = opcode.CustomCode().decode("utf-8")
op_set.add(opname)
else:
op_set.add(visualize.BuiltinCodeToName(builtin_code))
return op_set
def get_output_shapes(model_data):
"""Returns a list of output shapes in the tflite model data."""
model = schema_fb.Model.GetRootAsModel(model_data, 0)
output_shapes = []
for subgraph_idx in range(model.SubgraphsLength()):
subgraph = model.Subgraphs(subgraph_idx)
for output_idx in range(subgraph.OutputsLength()):
output_tensor_idx = subgraph.Outputs(output_idx)
output_tensor = subgraph.Tensors(output_tensor_idx)
output_shapes.append(output_tensor.ShapeAsNumpy().tolist())
return output_shapes
tensorflow/lite/python/test_util_test.py
+39
View File
@@ -0,0 +1,39 @@
# Copyright 2020 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.
# ==============================================================================
"""Tests for test_util.py."""
from tensorflow.lite.python import test_util as tflite_test_util
from tensorflow.python.framework import test_util
from tensorflow.python.platform import gfile
from tensorflow.python.platform import resource_loader
from tensorflow.python.platform import test
class TestUtilTest(test_util.TensorFlowTestCase):
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'])
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'])
if __name__ == '__main__':
test.main()
tensorflow/lite/python/testdata/BUILD
Vendored
+193
View File
@@ -0,0 +1,193 @@
load("@rules_cc//cc:cc_binary.bzl", "cc_binary")
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
load(
"//tensorflow:tensorflow.bzl",
"tf_custom_op_library",
"tf_gen_op_wrapper_py",
"tf_opts_nortti_if_android",
)
load("//tensorflow:tensorflow.default.bzl", "pybind_extension", "tf_custom_op_py_strict_library")
load("//tensorflow/lite:build_def.bzl", "tf_to_tflite", "tflite_copts")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
default_visibility = ["//tensorflow:internal"],
licenses = ["notice"],
)
exports_files(glob([
"*.pb",
"*.pbtxt",
]))
tf_to_tflite(
name = "permute_float",
src = "permute.pbtxt",
out = "permute_float.tflite",
options = [
"--input_arrays=input",
"--output_arrays=output",
],
)
tf_to_tflite(
name = "permute_uint8",
src = "permute.pbtxt",
out = "permute_uint8.tflite",
options = [
"--input_arrays=input",
"--output_arrays=output",
"--inference_type=QUANTIZED_UINT8",
"--std_dev_values=1",
"--mean_values=0",
"--default_ranges_min=0",
"--default_ranges_max=255",
],
)
tf_to_tflite(
name = "gather_string",
src = "gather.pbtxt",
out = "gather_string.tflite",
options = [
"--input_arrays=input,indices",
"--output_arrays=output",
],
)
tf_to_tflite(
name = "gather_string_0d",
src = "gather_0d.pbtxt",
out = "gather_string_0d.tflite",
options = [
"--input_arrays=input,indices",
"--output_arrays=output",
],
)
filegroup(
name = "interpreter_test_data",
srcs = [
"pc_conv.bin",
"two_signatures.tflite",
":gather_string",
":gather_string_0d",
":permute_float",
":permute_uint8",
],
visibility = ["//tensorflow:__subpackages__"],
)
cc_library(
name = "test_delegate",
testonly = 1,
srcs = ["test_delegate.cc"],
visibility = ["//visibility:private"],
deps = [
"//tensorflow/lite/core/c:common",
],
alwayslink = 1,
)
cc_binary(
name = "test_delegate.so",
testonly = 1,
linkshared = 1,
linkstatic = 1,
deps = [
":test_delegate",
],
)
cc_library(
name = "double_op_and_kernels",
testonly = 1,
srcs = ["double_op.cc"],
copts = tflite_copts() + tf_opts_nortti_if_android() + select({
"//tensorflow:android": ["-Wno-private-header"],
"//conditions:default": [],
}),
deps = select({
"//tensorflow:android": [
"//tensorflow/core:portable_tensorflow_lib_lite",
],
"//tensorflow:ios": [
"//tensorflow/core:portable_tensorflow_lib_lite",
],
"//conditions:default": [
"//tensorflow/core:framework",
"//tensorflow/core:lib",
],
}),
alwayslink = 1,
)
tf_custom_op_library(
name = "_double_op.so",
testonly = 1,
srcs = ["double_op.cc"],
)
tf_gen_op_wrapper_py(
name = "gen_double_op_wrapper",
testonly = 1,
out = "double_op_wrapper.py",
extra_py_deps = [
"//tensorflow/python:pywrap_tfe",
"//tensorflow/python/util:dispatch",
"//tensorflow/python/util:deprecation",
"//tensorflow/python/util:tf_export",
],
py_lib_rule = py_library,
deps = [":double_op_and_kernels"],
)
tf_custom_op_py_strict_library(
name = "double_op",
testonly = 1,
srcs = ["double_op.py"],
dso = [":_double_op.so"],
kernels = [":double_op_and_kernels"],
deps = [
":gen_double_op_wrapper",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:load_library",
"//tensorflow/python/platform:resource_loader",
],
)
cc_library(
name = "test_registerer",
srcs = ["test_registerer.cc"],
hdrs = ["test_registerer.h"],
visibility = ["//visibility:private"],
deps = [
"//tensorflow/lite:framework",
"//tensorflow/lite/kernels:builtin_ops",
"//tensorflow/lite/kernels:kernel_util",
"//tensorflow/lite/kernels/internal:tensor",
],
alwayslink = 1,
)
pybind_extension(
name = "_pywrap_test_registerer",
srcs = [
"test_registerer_wrapper.cc",
],
hdrs = ["test_registerer.h"],
additional_exported_symbols = ["TF_TestRegisterer"],
enable_stub_generation = True,
link_in_framework = True,
pytype_srcs = [
"_pywrap_test_registerer.pyi",
],
deps = [
":test_registerer",
"//tensorflow/lite:framework",
"//tensorflow/lite/kernels:builtin_ops",
"@pybind11",
"@xla//third_party/python_runtime:headers",
],
)
tensorflow/lite/python/testdata/_pywrap_test_registerer.pyi
Vendored
+17
View File
@@ -0,0 +1,17 @@
# Copyright 2023 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.
# ==============================================================================
def TF_TestRegisterer(arg0: int) -> None: ...
def get_num_test_registerer_calls() -> int: ...
tensorflow/lite/python/testdata/control_flow_v1.pbtxt
Vendored
+64
View File
@@ -0,0 +1,64 @@
node {
name: "a"
op: "Placeholder"
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
}
node {
name: "b"
op: "Placeholder"
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
}
node {
name: "c"
op: "Placeholder"
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
}
node {
name: "d"
op: "Placeholder"
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
}
node {
name: "Merge"
op: "Merge"
input: "a"
input: "b"
input: "c"
input: "d"
attr {
key: "N"
value {
i: 4
}
}
attr {
key: "T"
value {
type: DT_FLOAT
}
}
}
versions {
producer: 27
}
tensorflow/lite/python/testdata/control_flow_v1_saved_model/BUILD
Vendored
+9
View File
@@ -0,0 +1,9 @@
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
default_visibility = ["//tensorflow:internal"],
licenses = ["notice"],
)
exports_files([
"saved_model.pb",
])
tensorflow/lite/python/testdata/control_flow_v1_saved_model/saved_model.pb
Vendored
BIN
View File
Binary file not shown.
tensorflow/lite/python/testdata/double_op.cc
Vendored
+60
View File
@@ -0,0 +1,60 @@
/* Copyright 2020 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.
==============================================================================*/
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/shape_inference.h"
namespace tensorflow {
REGISTER_OP("Double")
.Input("input: T")
.Output("doubled: T")
.Attr("T: {int32, float}")
.SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
c->set_output(0, c->input(0));
return absl::OkStatus();
});
template <typename T>
class DoubleOp : public OpKernel {
public:
explicit DoubleOp(OpKernelConstruction* context) : OpKernel(context) {}
void Compute(OpKernelContext* context) override {
// Grab the input tensor
const Tensor& input_tensor = context->input(0);
auto input_flat = input_tensor.flat<T>();
// Create an output tensor
Tensor* output_tensor = nullptr;
OP_REQUIRES_OK(context, context->allocate_output(0, input_tensor.shape(),
&output_tensor));
auto output_flat = output_tensor->flat<T>();
// Set all but the first element of the output tensor to 0.
const int N = input_flat.size();
for (int i = 0; i < N; i++) {
output_flat(i) = 2 * input_flat(i);
}
}
};
REGISTER_KERNEL_BUILDER(
Name("Double").Device(DEVICE_CPU).TypeConstraint<int32_t>("T"),
DoubleOp<int32_t>);
REGISTER_KERNEL_BUILDER(
Name("Double").Device(DEVICE_CPU).TypeConstraint<float>("T"),
DoubleOp<float>);
} // namespace tensorflow
tensorflow/lite/python/testdata/double_op.py
Vendored
+31
View File
@@ -0,0 +1,31 @@
# Copyright 2020 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.
# ==============================================================================
"""Double op is a user's defined op for testing purpose."""
from tensorflow.lite.python.testdata import double_op_wrapper
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import load_library
from tensorflow.python.platform import resource_loader
_double_op = load_library.load_op_library(
resource_loader.get_path_to_datafile('_double_op.so'))
def double(input_tensor):
"""Double op applies element-wise double to input data."""
if (input_tensor.dtype != dtypes.int32 and
input_tensor.dtype != dtypes.float32):
raise ValueError('Double op only accept int32 or float32 values.')
return double_op_wrapper.double(input_tensor)
tensorflow/lite/python/testdata/gather.pbtxt
Vendored
+93
View File
@@ -0,0 +1,93 @@
node {
name: "input"
op: "Placeholder"
device: "/device:CPU:0"
attr {
key: "dtype"
value {
type: DT_STRING
}
}
attr {
key: "shape"
value {
shape {
dim {
size: 10
}
}
}
}
}
node {
name: "indices"
op: "Placeholder"
device: "/device:CPU:0"
attr {
key: "dtype"
value {
type: DT_INT64
}
}
attr {
key: "shape"
value {
shape {
dim {
size: 3
}
}
}
}
}
node {
name: "axis"
op: "Const"
device: "/device:CPU:0"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 0
}
}
}
}
node {
name: "output"
op: "GatherV2"
input: "input"
input: "indices"
input: "axis"
device: "/device:CPU:0"
attr {
key: "Taxis"
value {
type: DT_INT32
}
}
attr {
key: "Tindices"
value {
type: DT_INT64
}
}
attr {
key: "Tparams"
value {
type: DT_STRING
}
}
}
versions {
producer: 27
}
tensorflow/lite/python/testdata/gather_0d.pbtxt
Vendored
+108
View File
@@ -0,0 +1,108 @@
node {
name: "input"
op: "Placeholder"
device: "/device:CPU:0"
attr {
key: "dtype"
value {
type: DT_STRING
}
}
}
node {
name: "input_const"
op: "Const"
device: "/device:CPU:0"
attr {
key: "dtype"
value {
type: DT_STRING
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_STRING
tensor_shape {
dim {
size: 1
}
}
string_val: "abcd"
}
}
}
}
node {
name: "indices"
op: "Placeholder"
device: "/device:CPU:0"
attr {
key: "dtype"
value {
type: DT_INT64
}
}
attr {
key: "shape"
value {
shape {
dim {
size: 3
}
}
}
}
}
node {
name: "axis"
op: "Const"
device: "/device:CPU:0"
attr {
key: "dtype"
value {
type: DT_INT32
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_INT32
tensor_shape {
}
int_val: 0
}
}
}
}
node {
name: "output"
op: "GatherV2"
input: "input_const"
input: "indices"
input: "axis"
device: "/device:CPU:0"
attr {
key: "Taxis"
value {
type: DT_INT32
}
}
attr {
key: "Tindices"
value {
type: DT_INT64
}
}
attr {
key: "Tparams"
value {
type: DT_STRING
}
}
}
versions {
producer: 27
}
tensorflow/lite/python/testdata/pc_conv.bin
Vendored
BIN
View File
Binary file not shown.
tensorflow/lite/python/testdata/permute.pbtxt
Vendored
+98
View File
@@ -0,0 +1,98 @@
node {
name: "input"
op: "Placeholder"
device: "/device:CPU:0"
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
attr {
key: "shape"
value {
shape {
dim {
size: 1
}
dim {
size: 4
}
}
}
}
}
node {
name: "Const"
op: "Const"
device: "/device:CPU:0"
attr {
key: "dtype"
value {
type: DT_FLOAT
}
}
attr {
key: "value"
value {
tensor {
dtype: DT_FLOAT
tensor_shape {
dim {
size: 4
}
dim {
size: 4
}
}
float_val: 0.0
float_val: 0.0
float_val: 0.0
float_val: 1.0
float_val: 0.0
float_val: 0.0
float_val: 1.0
float_val: 0.0
float_val: 0.0
float_val: 1.0
float_val: 0.0
float_val: 0.0
float_val: 1.0
float_val: 0.0
float_val: 0.0
float_val: 0.0
}
}
}
}
node {
name: "output"
op: "MatMul"
input: "input"
input: "Const"
device: "/device:CPU:0"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "transpose_a"
value {
b: false
}
}
attr {
key: "transpose_b"
value {
b: false
}
}
}
versions {
producer: 27
}
tensorflow/lite/python/testdata/test_delegate.cc
Vendored
+94
View File
@@ -0,0 +1,94 @@
/* 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.
==============================================================================*/
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include "tensorflow/lite/core/c/common.h"
namespace tflite {
namespace {
int num_delegates_created = 0;
int num_delegates_destroyed = 0;
int num_delegates_invoked = 0;
int options_counter = 0;
int (*destruction_callback)(const char* s) = nullptr;
typedef void (*ErrorHandler)(const char*);
} // namespace
extern "C" {
TfLiteDelegate* tflite_plugin_create_delegate(char** options_keys,
char** options_values,
size_t num_options,
ErrorHandler error_handler) {
num_delegates_created++;
for (int idx = 0; idx < num_options; idx++) {
if (std::strcmp("options_counter", options_keys[idx]) == 0) {
int int_value;
if (sscanf(options_values[idx], "%d", &int_value) == 1) {
options_counter += int_value;
}
} else if (std::strcmp("fail", options_keys[idx]) == 0) {
if (error_handler) error_handler("Fail argument sent.");
return nullptr;
}
}
TfLiteDelegate* ptr = new TfLiteDelegate;
*ptr = TfLiteDelegateCreate();
ptr->Prepare = [](TfLiteContext* context, TfLiteDelegate* delegate) {
num_delegates_invoked++;
return kTfLiteOk;
};
ptr->flags = kTfLiteDelegateFlagsNone;
return ptr;
}
void set_destroy_callback(int (*callback)(const char* s)) {
destruction_callback = callback;
}
void tflite_plugin_destroy_delegate(TfLiteDelegate* delegate) {
num_delegates_destroyed++;
delete delegate;
if (destruction_callback) {
destruction_callback("test_delegate");
// destruction_callback is a global variable,
// so it should be set to nullptr here to avoid crashes
destruction_callback = nullptr;
}
}
void initialize_counters() {
num_delegates_created = 0;
num_delegates_destroyed = 0;
num_delegates_invoked = 0;
options_counter = 0;
}
int get_num_delegates_created() { return num_delegates_created; }
int get_num_delegates_destroyed() { return num_delegates_destroyed; }
int get_num_delegates_invoked() { return num_delegates_invoked; }
int get_options_counter() { return options_counter; }
}
} // namespace tflite
tensorflow/lite/python/testdata/test_registerer.cc
Vendored
+107
View File
@@ -0,0 +1,107 @@
/* 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.
==============================================================================*/
#include "tensorflow/lite/python/testdata/test_registerer.h"
#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
#include "tensorflow/lite/kernels/kernel_util.h"
namespace tflite {
namespace {
static int num_test_registerer_calls = 0;
TfLiteRegistration* GetFakeRegistration() {
static TfLiteRegistration fake_op;
return &fake_op;
}
namespace double_op {
constexpr int kInputTensor = 0;
constexpr int kOutputTensor = 0;
TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
TF_LITE_ENSURE_EQ(context, NumInputs(node), 1);
TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
const TfLiteTensor* input;
TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
TfLiteTensor* output;
TF_LITE_ENSURE_OK(context,
GetOutputSafe(context, node, kOutputTensor, &output));
TfLiteIntArray* output_shape = TfLiteIntArrayCopy(input->dims);
TF_LITE_ENSURE_TYPES_EQ(context, output->type, input->type);
return context->ResizeTensor(context, output, output_shape);
}
TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
const TfLiteTensor* input;
TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, kInputTensor, &input));
TfLiteTensor* output;
TF_LITE_ENSURE_OK(context,
GetOutputSafe(context, node, kOutputTensor, &output));
TF_LITE_ENSURE_TYPES_EQ(context, output->type, input->type);
const size_t size = GetTensorShape(input).FlatSize();
if (input->type == kTfLiteFloat32) {
const float* input_ptr = input->data.f;
float* output_ptr = output->data.f;
for (int i = 0; i < size; ++i) {
output_ptr[i] = input_ptr[i] + input_ptr[i];
}
} else if (input->type == kTfLiteInt32) {
const int32_t* input_ptr = input->data.i32;
int32_t* output_ptr = output->data.i32;
for (int i = 0; i < size; ++i) {
output_ptr[i] = input_ptr[i] + input_ptr[i];
}
} else {
return kTfLiteError;
}
return kTfLiteOk;
}
} // namespace double_op
TfLiteRegistration* GetDoubleRegistration() {
static TfLiteRegistration double_op = {nullptr, nullptr, double_op::Prepare,
double_op::Eval};
return &double_op;
}
} // namespace
// Dummy registerer function with the correct signature. Registers a fake custom
// op needed by test models. Increments the num_test_registerer_calls counter by
// one. The TF_ prefix is needed to get past the version script in the OSS
// build.
extern "C" void TF_TestRegisterer(tflite::MutableOpResolver *resolver) {
resolver->AddCustom("FakeOp", GetFakeRegistration());
resolver->AddCustom("Double", GetDoubleRegistration());
num_test_registerer_calls++;
}
// Returns the num_test_registerer_calls counter and re-sets it.
int get_num_test_registerer_calls() {
const int result = num_test_registerer_calls;
num_test_registerer_calls = 0;
return result;
}
} // namespace tflite
tensorflow/lite/python/testdata/test_registerer.h
Vendored
+32
View File
@@ -0,0 +1,32 @@
/* 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.
==============================================================================*/
#ifndef TENSORFLOW_LITE_PYTHON_TESTDATA_TEST_REGISTERER_H_
#define TENSORFLOW_LITE_PYTHON_TESTDATA_TEST_REGISTERER_H_
#include "tensorflow/lite/mutable_op_resolver.h"
namespace tflite {
// Dummy registerer function with the correct signature. Ignores the resolver
// but increments the num_test_registerer_calls counter by one. The TF_ prefix
// is needed to get past the version script in the OSS build.
extern "C" void TF_TestRegisterer(tflite::MutableOpResolver *resolver);
// Returns the num_test_registerer_calls counter and re-sets it.
int get_num_test_registerer_calls();
} // namespace tflite
#endif // TENSORFLOW_LITE_PYTHON_TESTDATA_TEST_REGISTERER_H_
tensorflow/lite/python/testdata/test_registerer_wrapper.cc
Vendored
+36
View File
@@ -0,0 +1,36 @@
/* 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.
==============================================================================*/
#include "pybind11/pybind11.h" // from @pybind11
#include "pybind11/pytypes.h" // from @pybind11
#include "tensorflow/lite/python/testdata/test_registerer.h"
PYBIND11_MODULE(_pywrap_test_registerer, m) {
m.doc() = R"pbdoc(
_pywrap_test_registerer
-----
)pbdoc";
m.def("get_num_test_registerer_calls", &tflite::get_num_test_registerer_calls,
R"pbdoc(
Returns the num_test_registerer_calls counter and re-sets it.
)pbdoc");
m.def(
"TF_TestRegisterer",
[](uintptr_t resolver) {
tflite::TF_TestRegisterer(
reinterpret_cast<tflite::MutableOpResolver*>(resolver));
},
R"pbdoc(
Dummy registerer function with the correct signature. Registers a fake
custom op needed by test models. Increments the
num_test_registerer_calls counter by one.
)pbdoc");
}
tensorflow/lite/python/tflite_convert.py
+696
View File
@@ -0,0 +1,696 @@
# Copyright 2018 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.
# ==============================================================================
"""Python command line interface for converting TF models to TFLite models."""
import argparse
import os
import sys
import warnings
from absl import app
import tensorflow as tf
from tensorflow.lite.python import lite
from tensorflow.lite.python.convert import register_custom_opdefs
from tensorflow.lite.toco import toco_flags_pb2 as _toco_flags_pb2
from tensorflow.lite.toco.logging import gen_html
from tensorflow.python import tf2
from tensorflow.python.framework import dtypes
from tensorflow.python.platform import gfile
from tensorflow.python.util import keras_deps
# Needed to enable TF2 by default.
_ = tf.keras.models.save_model # ensure necessary imports are executed
def _parse_array(values, type_fn=str):
if values is not None:
return [type_fn(val) for val in values.split(",") if val]
return None
def _parse_set(values):
if values is not None:
return set([item for item in values.split(",") if item])
return None
def _parse_inference_type(value, flag):
"""Converts the inference type to the value of the constant.
Args:
value: str representing the inference type.
flag: str representing the flag name.
Returns:
tf.dtype.
Raises:
ValueError: Unsupported value.
"""
if value == "FLOAT":
return dtypes.float32
if value == "INT8":
return dtypes.int8
if value == "UINT8" or value == "QUANTIZED_UINT8":
return dtypes.uint8
raise ValueError(
"Unsupported value for `{}` flag. Expected FLOAT, INT8, UINT8, or "
"QUANTIZED_UINT8 instead got {}.".format(flag, value))
class _ParseBooleanFlag(argparse.Action):
"""Helper class to parse boolean flag that optionally accepts truth value."""
def __init__(self, option_strings, dest, nargs=None, **kwargs):
if nargs != "?":
# This should never happen. This class is only used once below with
# nargs="?".
raise ValueError(
"This parser only supports nargs='?' (0 or 1 additional arguments)")
super(_ParseBooleanFlag, self).__init__(
option_strings, dest, nargs=nargs, **kwargs)
def __call__(self, parser, namespace, values, option_string=None):
if values is None:
# Handling `--boolean_flag`.
# Without additional arguments, it implies true.
flag_value = True
elif values.lower() == "true":
# Handling `--boolean_flag=true`.
# (Case insensitive after the equal sign)
flag_value = True
elif values.lower() == "false":
# Handling `--boolean_flag=false`.
# (Case insensitive after the equal sign)
flag_value = False
else:
raise ValueError("Invalid argument to --{}. Must use flag alone,"
" or specify true/false.".format(self.dest))
setattr(namespace, self.dest, flag_value)
def _get_tflite_converter(flags):
"""Makes a TFLiteConverter object based on the flags provided.
Args:
flags: argparse.Namespace object containing TFLite flags.
Returns:
TFLiteConverter object.
Raises:
ValueError: Invalid flags.
"""
# Parse input and output arrays.
input_arrays = _parse_array(flags.input_arrays)
input_shapes = None
if flags.input_shapes:
input_shapes_list = [
_parse_array(shape, type_fn=int)
for shape in flags.input_shapes.split(":")
]
input_shapes = dict(list(zip(input_arrays, input_shapes_list)))
output_arrays = _parse_array(flags.output_arrays)
converter_kwargs = {
"input_arrays": input_arrays,
"input_shapes": input_shapes,
"output_arrays": output_arrays
}
# Create TFLiteConverter.
if flags.graph_def_file:
converter_fn = lite.TFLiteConverter.from_frozen_graph
converter_kwargs["graph_def_file"] = flags.graph_def_file
elif flags.saved_model_dir:
converter_fn = lite.TFLiteConverter.from_saved_model
converter_kwargs["saved_model_dir"] = flags.saved_model_dir
converter_kwargs["tag_set"] = _parse_set(flags.saved_model_tag_set)
converter_kwargs["signature_key"] = flags.saved_model_signature_key
elif flags.keras_model_file:
converter_fn = lite.TFLiteConverter.from_keras_model_file
converter_kwargs["model_file"] = flags.keras_model_file
else:
raise ValueError("--graph_def_file, --saved_model_dir, or "
"--keras_model_file must be specified.")
return converter_fn(**converter_kwargs)
def _convert_tf1_model(flags):
"""Calls function to convert the TensorFlow 1.X model into a TFLite model.
Args:
flags: argparse.Namespace object.
Raises:
ValueError: Invalid flags.
"""
# Register custom opdefs before converter object creation.
if flags.custom_opdefs:
register_custom_opdefs(_parse_array(flags.custom_opdefs))
# Create converter.
converter = _get_tflite_converter(flags)
if flags.inference_type:
converter.inference_type = _parse_inference_type(flags.inference_type,
"inference_type")
if flags.inference_input_type:
converter.inference_input_type = _parse_inference_type(
flags.inference_input_type, "inference_input_type")
if flags.output_format:
converter.output_format = _toco_flags_pb2.FileFormat.Value(
flags.output_format)
if flags.mean_values and flags.std_dev_values:
input_arrays = converter.get_input_arrays()
std_dev_values = _parse_array(flags.std_dev_values, type_fn=float)
# In quantized inference, mean_value has to be integer so that the real
# value 0.0 is exactly representable.
if converter.inference_type == dtypes.float32:
mean_values = _parse_array(flags.mean_values, type_fn=float)
else:
mean_values = _parse_array(flags.mean_values, type_fn=int)
quant_stats = list(zip(mean_values, std_dev_values))
if ((not flags.input_arrays and len(input_arrays) > 1) or
(len(input_arrays) != len(quant_stats))):
raise ValueError("Mismatching --input_arrays, --std_dev_values, and "
"--mean_values. The flags must have the same number of "
"items. The current input arrays are '{0}'. "
"--input_arrays must be present when specifying "
"--std_dev_values and --mean_values with multiple input "
"tensors in order to map between names and "
"values.".format(",".join(input_arrays)))
converter.quantized_input_stats = dict(list(zip(input_arrays, quant_stats)))
if (flags.default_ranges_min is not None) and (flags.default_ranges_max is
not None):
converter.default_ranges_stats = (flags.default_ranges_min,
flags.default_ranges_max)
if flags.drop_control_dependency:
converter.drop_control_dependency = flags.drop_control_dependency
if flags.reorder_across_fake_quant:
converter.reorder_across_fake_quant = flags.reorder_across_fake_quant
if flags.change_concat_input_ranges:
converter.change_concat_input_ranges = (
flags.change_concat_input_ranges == "TRUE")
if flags.allow_custom_ops:
converter.allow_custom_ops = flags.allow_custom_ops
if flags.target_ops:
ops_set_options = lite.OpsSet.get_options()
converter.target_spec.supported_ops = set()
for option in flags.target_ops.split(","):
if option not in ops_set_options:
raise ValueError("Invalid value for --target_ops. Options: "
"{0}".format(",".join(ops_set_options)))
converter.target_spec.supported_ops.add(lite.OpsSet(option))
if flags.experimental_select_user_tf_ops:
if lite.OpsSet.SELECT_TF_OPS not in converter.target_spec.supported_ops:
raise ValueError("--experimental_select_user_tf_ops can only be set if "
"--target_ops contains SELECT_TF_OPS.")
user_op_set = set()
for op_name in flags.experimental_select_user_tf_ops.split(","):
user_op_set.add(op_name)
converter.target_spec.experimental_select_user_tf_ops = list(user_op_set)
if flags.post_training_quantize:
converter.optimizations = [lite.Optimize.DEFAULT]
if converter.inference_type != dtypes.float32:
print("--post_training_quantize quantizes a graph of inference_type "
"FLOAT. Overriding inference_type to FLOAT.")
converter.inference_type = dtypes.float32
if flags.quantize_to_float16:
converter.target_spec.supported_types = [dtypes.float16]
if not flags.post_training_quantize:
print("--quantize_to_float16 will only take effect with the "
"--post_training_quantize flag enabled.")
if flags.dump_graphviz_dir:
converter.dump_graphviz_dir = flags.dump_graphviz_dir
if flags.dump_graphviz_video:
converter.dump_graphviz_vode = flags.dump_graphviz_video
if flags.conversion_summary_dir:
converter.conversion_summary_dir = flags.conversion_summary_dir
converter.experimental_new_converter = flags.experimental_new_converter
if flags.experimental_new_quantizer is not None:
converter.experimental_new_quantizer = flags.experimental_new_quantizer
# Convert model.
output_data = converter.convert()
with gfile.GFile(flags.output_file, "wb") as f:
f.write(output_data)
def _convert_tf2_model(flags):
"""Calls function to convert the TensorFlow 2.0 model into a TFLite model.
Args:
flags: argparse.Namespace object.
Raises:
ValueError: Unsupported file format.
"""
# Load the model.
if flags.saved_model_dir:
converter = lite.TFLiteConverterV2.from_saved_model(
flags.saved_model_dir,
signature_keys=_parse_array(flags.saved_model_signature_key),
tags=_parse_set(flags.saved_model_tag_set))
elif flags.keras_model_file:
model = keras_deps.get_load_model_function()(flags.keras_model_file)
converter = lite.TFLiteConverterV2.from_keras_model(model)
converter.experimental_new_converter = flags.experimental_new_converter
if flags.experimental_new_quantizer is not None:
converter.experimental_new_quantizer = flags.experimental_new_quantizer
# Convert the model.
tflite_model = converter.convert()
with gfile.GFile(flags.output_file, "wb") as f:
f.write(tflite_model)
def _check_tf1_flags(flags, unparsed):
"""Checks the parsed and unparsed flags to ensure they are valid in 1.X.
Raises an error if previously support unparsed flags are found. Raises an
error for parsed flags that don't meet the required conditions.
Args:
flags: argparse.Namespace object containing TFLite flags.
unparsed: List of unparsed flags.
Raises:
ValueError: Invalid flags.
"""
# Check unparsed flags for common mistakes based on previous TOCO.
def _get_message_unparsed(flag, orig_flag, new_flag):
if flag.startswith(orig_flag):
return "\n Use {0} instead of {1}".format(new_flag, orig_flag)
return ""
if unparsed:
output = ""
for flag in unparsed:
output += _get_message_unparsed(flag, "--input_file", "--graph_def_file")
output += _get_message_unparsed(flag, "--savedmodel_directory",
"--saved_model_dir")
output += _get_message_unparsed(flag, "--std_value", "--std_dev_values")
output += _get_message_unparsed(flag, "--batch_size", "--input_shapes")
output += _get_message_unparsed(flag, "--dump_graphviz",
"--dump_graphviz_dir")
if output:
raise ValueError(output)
# Check that flags are valid.
if flags.graph_def_file and (not flags.input_arrays or
not flags.output_arrays):
raise ValueError("--input_arrays and --output_arrays are required with "
"--graph_def_file")
if flags.input_shapes:
if not flags.input_arrays:
raise ValueError("--input_shapes must be used with --input_arrays")
if flags.input_shapes.count(":") != flags.input_arrays.count(","):
raise ValueError("--input_shapes and --input_arrays must have the same "
"number of items")
if flags.std_dev_values or flags.mean_values:
if bool(flags.std_dev_values) != bool(flags.mean_values):
raise ValueError("--std_dev_values and --mean_values must be used "
"together")
if flags.std_dev_values.count(",") != flags.mean_values.count(","):
raise ValueError("--std_dev_values, --mean_values must have the same "
"number of items")
if (flags.default_ranges_min is None) != (flags.default_ranges_max is None):
raise ValueError("--default_ranges_min and --default_ranges_max must be "
"used together")
if flags.dump_graphviz_video and not flags.dump_graphviz_dir:
raise ValueError("--dump_graphviz_video must be used with "
"--dump_graphviz_dir")
if flags.custom_opdefs and not flags.experimental_new_converter:
raise ValueError("--custom_opdefs must be used with "
"--experimental_new_converter")
if flags.custom_opdefs and not flags.allow_custom_ops:
raise ValueError("--custom_opdefs must be used with --allow_custom_ops")
if (flags.experimental_select_user_tf_ops and
not flags.experimental_new_converter):
raise ValueError("--experimental_select_user_tf_ops must be used with "
"--experimental_new_converter")
def _check_tf2_flags(flags):
"""Checks the parsed and unparsed flags to ensure they are valid in 2.X.
Args:
flags: argparse.Namespace object containing TFLite flags.
Raises:
ValueError: Invalid flags.
"""
if not flags.keras_model_file and not flags.saved_model_dir:
raise ValueError("one of the arguments --saved_model_dir "
"--keras_model_file is required")
def _get_tf1_flags(parser):
"""Returns ArgumentParser for tflite_convert for TensorFlow 1.X.
Args:
parser: ArgumentParser
"""
# Input file flags.
input_file_group = parser.add_mutually_exclusive_group(required=True)
input_file_group.add_argument(
"--graph_def_file",
type=str,
help="Full filepath of file containing frozen TensorFlow GraphDef.")
input_file_group.add_argument(
"--saved_model_dir",
type=str,
help="Full filepath of directory containing the SavedModel.")
input_file_group.add_argument(
"--keras_model_file",
type=str,
help="Full filepath of HDF5 file containing tf.Keras model.")
# Model format flags.
parser.add_argument(
"--output_format",
type=str.upper,
choices=["TFLITE", "GRAPHVIZ_DOT"],
help="Output file format.")
parser.add_argument(
"--inference_type",
type=str.upper,
default="FLOAT",
help=("Target data type of real-number arrays in the output file. "
"Must be either FLOAT, INT8 or UINT8."))
parser.add_argument(
"--inference_input_type",
type=str.upper,
help=("Target data type of real-number input arrays. Allows for a "
"different type for input arrays in the case of quantization. "
"Must be either FLOAT, INT8 or UINT8."))
# Input and output arrays flags.
parser.add_argument(
"--input_arrays",
type=str,
help="Names of the input arrays, comma-separated.")
parser.add_argument(
"--input_shapes",
type=str,
help="Shapes corresponding to --input_arrays, colon-separated.")
parser.add_argument(
"--output_arrays",
type=str,
help="Names of the output arrays, comma-separated.")
# SavedModel related flags.
parser.add_argument(
"--saved_model_tag_set",
type=str,
help=("Comma-separated set of tags identifying the MetaGraphDef within "
"the SavedModel to analyze. All tags must be present. In order to "
"pass in an empty tag set, pass in \"\". (default \"serve\")"))
parser.add_argument(
"--saved_model_signature_key",
type=str,
help=("Key identifying the SignatureDef containing inputs and outputs. "
"(default DEFAULT_SERVING_SIGNATURE_DEF_KEY)"))
# Quantization flags.
parser.add_argument(
"--std_dev_values",
type=str,
help=("Standard deviation of training data for each input tensor, "
"comma-separated floats. Used for quantized input tensors. "
"(default None)"))
parser.add_argument(
"--mean_values",
type=str,
help=("Mean of training data for each input tensor, comma-separated "
"floats. Used for quantized input tensors. (default None)"))
parser.add_argument(
"--default_ranges_min",
type=float,
help=("Default value for min bound of min/max range values used for all "
"arrays without a specified range, Intended for experimenting with "
"quantization via \"dummy quantization\". (default None)"))
parser.add_argument(
"--default_ranges_max",
type=float,
help=("Default value for max bound of min/max range values used for all "
"arrays without a specified range, Intended for experimenting with "
"quantization via \"dummy quantization\". (default None)"))
# quantize_weights is DEPRECATED.
parser.add_argument(
"--quantize_weights",
dest="post_training_quantize",
action="store_true",
help=argparse.SUPPRESS)
parser.add_argument(
"--post_training_quantize",
dest="post_training_quantize",
action="store_true",
help=(
"Boolean indicating whether to quantize the weights of the "
"converted float model. Model size will be reduced and there will "
"be latency improvements (at the cost of accuracy). (default False)"))
parser.add_argument(
"--quantize_to_float16",
dest="quantize_to_float16",
action="store_true",
help=("Boolean indicating whether to quantize weights to fp16 instead of "
"the default int8 when post-training quantization "
"(--post_training_quantize) is enabled. (default False)"))
# Graph manipulation flags.
parser.add_argument(
"--drop_control_dependency",
action="store_true",
help=("Boolean indicating whether to drop control dependencies silently. "
"This is due to TensorFlow not supporting control dependencies. "
"(default True)"))
parser.add_argument(
"--reorder_across_fake_quant",
action="store_true",
help=("Boolean indicating whether to reorder FakeQuant nodes in "
"unexpected locations. Used when the location of the FakeQuant "
"nodes is preventing graph transformations necessary to convert "
"the graph. Results in a graph that differs from the quantized "
"training graph, potentially causing differing arithmetic "
"behavior. (default False)"))
# Usage for this flag is --change_concat_input_ranges=true or
# --change_concat_input_ranges=false in order to make it clear what the flag
# is set to. This keeps the usage consistent with other usages of the flag
# where the default is different. The default value here is False.
parser.add_argument(
"--change_concat_input_ranges",
type=str.upper,
choices=["TRUE", "FALSE"],
help=("Boolean to change behavior of min/max ranges for inputs and "
"outputs of the concat operator for quantized models. Changes the "
"ranges of concat operator overlap when true. (default False)"))
# Permitted ops flags.
parser.add_argument(
"--allow_custom_ops",
action=_ParseBooleanFlag,
nargs="?",
help=("Boolean indicating whether to allow custom operations. When false "
"any unknown operation is an error. When true, custom ops are "
"created for any op that is unknown. The developer will need to "
"provide these to the TensorFlow Lite runtime with a custom "
"resolver. (default False)"))
parser.add_argument(
"--custom_opdefs",
type=str,
help=("String representing a list of custom ops OpDefs delineated with "
"commas that are included in the GraphDef. Required when using "
"custom operations with --experimental_new_converter."))
parser.add_argument(
"--target_ops",
type=str,
help=("Experimental flag, subject to change. Set of OpsSet options "
"indicating which converter to use. Options: {0}. One or more "
"option may be specified. (default set([OpsSet.TFLITE_BUILTINS]))"
"".format(",".join(lite.OpsSet.get_options()))))
parser.add_argument(
"--experimental_select_user_tf_ops",
type=str,
help=("Experimental flag, subject to change. Comma separated list of "
"user's defined TensorFlow operators required in the runtime."))
# Logging flags.
parser.add_argument(
"--dump_graphviz_dir",
type=str,
help=("Full filepath of folder to dump the graphs at various stages of "
"processing GraphViz .dot files. Preferred over --output_format="
"GRAPHVIZ_DOT in order to keep the requirements of the output "
"file."))
parser.add_argument(
"--dump_graphviz_video",
action="store_true",
help=("Boolean indicating whether to dump the graph after every graph "
"transformation"))
parser.add_argument(
"--conversion_summary_dir",
type=str,
help=("Full filepath to store the conversion logs, which includes "
"graphviz of the model before/after the conversion, an HTML report "
"and the conversion proto buffers. This will only be generated "
"when passing --experimental_new_converter"))
def _get_tf2_flags(parser):
"""Returns ArgumentParser for tflite_convert for TensorFlow 2.0.
Args:
parser: ArgumentParser
"""
# Input file flags.
input_file_group = parser.add_mutually_exclusive_group()
input_file_group.add_argument(
"--saved_model_dir",
type=str,
help="Full path of the directory containing the SavedModel.")
input_file_group.add_argument(
"--keras_model_file",
type=str,
help="Full filepath of HDF5 file containing tf.Keras model.")
# SavedModel related flags.
parser.add_argument(
"--saved_model_tag_set",
type=str,
help=("Comma-separated set of tags identifying the MetaGraphDef within "
"the SavedModel to analyze. All tags must be present. In order to "
"pass in an empty tag set, pass in \"\". (default \"serve\")"))
parser.add_argument(
"--saved_model_signature_key",
type=str,
help=("Key identifying the SignatureDef containing inputs and outputs. "
"(default DEFAULT_SERVING_SIGNATURE_DEF_KEY)"))
# Enables 1.X converter in 2.X.
parser.add_argument(
"--enable_v1_converter",
action="store_true",
help=("Enables the TensorFlow V1 converter in 2.0"))
def _get_parser(use_v2_converter):
"""Returns an ArgumentParser for tflite_convert.
Args:
use_v2_converter: Indicates which converter to return.
Return: ArgumentParser.
"""
parser = argparse.ArgumentParser(
description=("Command line tool to run TensorFlow Lite Converter."))
# Output file flag.
parser.add_argument(
"--output_file",
type=str,
help="Full filepath of the output file.",
required=True)
if use_v2_converter:
_get_tf2_flags(parser)
else:
_get_tf1_flags(parser)
parser.add_argument(
"--experimental_new_converter",
action=_ParseBooleanFlag,
nargs="?",
default=True,
help=("Experimental flag, subject to change. Enables MLIR-based "
"conversion instead of TOCO conversion. (default True)"))
parser.add_argument(
"--experimental_new_quantizer",
action=_ParseBooleanFlag,
nargs="?",
help=("Experimental flag, subject to change. Enables MLIR-based "
"quantizer instead of flatbuffer conversion. (default True)"))
return parser
def run_main(_):
"""Main in tflite_convert.py."""
use_v2_converter = tf2.enabled()
parser = _get_parser(use_v2_converter)
tflite_flags, unparsed = parser.parse_known_args(args=sys.argv[1:])
# If the user is running TensorFlow 2.X but has passed in enable_v1_converter
# then parse the flags again with the 1.X converter flags.
if tf2.enabled() and tflite_flags.enable_v1_converter:
use_v2_converter = False
parser = _get_parser(use_v2_converter)
tflite_flags, unparsed = parser.parse_known_args(args=sys.argv[1:])
# Checks if the flags are valid.
try:
if use_v2_converter:
_check_tf2_flags(tflite_flags)
else:
_check_tf1_flags(tflite_flags, unparsed)
except ValueError as e:
parser.print_usage()
file_name = os.path.basename(sys.argv[0])
sys.stderr.write("{0}: error: {1}\n".format(file_name, str(e)))
sys.exit(1)
# Convert the model according to the user provided flag.
if use_v2_converter:
_convert_tf2_model(tflite_flags)
else:
try:
_convert_tf1_model(tflite_flags)
finally:
if tflite_flags.conversion_summary_dir:
if tflite_flags.experimental_new_converter:
gen_html.gen_conversion_log_html(tflite_flags.conversion_summary_dir,
tflite_flags.post_training_quantize,
tflite_flags.output_file)
else:
warnings.warn(
"Conversion summary will only be generated when enabling"
" the new converter via --experimental_new_converter. ")
def main():
app.run(main=run_main, argv=sys.argv[:1])
if __name__ == "__main__":
main()
tensorflow/lite/python/tflite_convert_test.py
+570
View File
@@ -0,0 +1,570 @@
# 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.
# ==============================================================================
"""Tests for tflite_convert.py."""
import os
from absl.testing import parameterized
import numpy as np
from tensorflow import keras
from tensorflow.core.framework import graph_pb2
from tensorflow.lite.python import test_util as tflite_test_util
from tensorflow.lite.python import tflite_convert
from tensorflow.lite.python.convert import register_custom_opdefs
from tensorflow.python import tf2
from tensorflow.python.client import session
from tensorflow.python.eager import def_function
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import test_util
from tensorflow.python.framework.importer import import_graph_def
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import random_ops
from tensorflow.python.platform import gfile
from tensorflow.python.platform import resource_loader
from tensorflow.python.platform import test
from tensorflow.python.saved_model import saved_model
from tensorflow.python.saved_model.save import save
from tensorflow.python.trackable import autotrackable
from tensorflow.python.training.training_util import write_graph
class TestModels(test_util.TensorFlowTestCase):
def _getFilepath(self, filename):
return os.path.join(self.get_temp_dir(), filename)
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.par')
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 _getKerasModelFile(self):
x = np.array([[1.], [2.]])
y = np.array([[2.], [4.]])
model = keras.models.Sequential([
keras.layers.Dropout(0.2, input_shape=(1,)),
keras.layers.Dense(1),
])
model.compile(optimizer='sgd', loss='mean_squared_error')
model.fit(x, y, epochs=1)
keras_file = self._getFilepath('model.h5')
keras.models.save_model(model, keras_file)
return keras_file
def _getKerasFunctionalModelFile(self):
"""Returns a functional Keras model with output shapes [[1, 1], [1, 2]]."""
input_tensor = keras.layers.Input(shape=(1,))
output1 = keras.layers.Dense(1, name='b')(input_tensor)
output2 = keras.layers.Dense(2, name='a')(input_tensor)
model = keras.models.Model(inputs=input_tensor, outputs=[output1, output2])
keras_file = self._getFilepath('functional_model.h5')
keras.models.save_model(model, keras_file)
return keras_file
class TfLiteConvertV1Test(TestModels):
def _run(self,
flags_str,
should_succeed,
expected_ops_in_converted_model=None):
if tf2.enabled():
flags_str += ' --enable_v1_converter'
super(TfLiteConvertV1Test, self)._run(flags_str, should_succeed,
expected_ops_in_converted_model)
def testFrozenGraphDef(self):
with ops.Graph().as_default():
in_tensor = array_ops.placeholder(
shape=[1, 16, 16, 3], dtype=dtypes.float32)
_ = in_tensor + in_tensor
sess = session.Session()
# Write graph to file.
graph_def_file = self._getFilepath('model.pb')
write_graph(sess.graph_def, '', graph_def_file, False)
sess.close()
flags_str = ('--graph_def_file={0} --input_arrays={1} '
'--output_arrays={2}'.format(graph_def_file, 'Placeholder',
'add'))
self._run(flags_str, should_succeed=True)
os.remove(graph_def_file)
# Run `tflite_convert` explicitly with the legacy converter.
# Before the new converter is enabled by default, this flag has no real
# effects.
def testFrozenGraphDefWithLegacyConverter(self):
with ops.Graph().as_default():
in_tensor = array_ops.placeholder(
shape=[1, 16, 16, 3], dtype=dtypes.float32)
_ = in_tensor + in_tensor
sess = session.Session()
# Write graph to file.
graph_def_file = self._getFilepath('model.pb')
write_graph(sess.graph_def, '', graph_def_file, False)
sess.close()
flags_str = (
'--graph_def_file={0} --input_arrays={1} '
'--output_arrays={2} --experimental_new_converter=false'.format(
graph_def_file, 'Placeholder', 'add'))
self._run(flags_str, should_succeed=True)
os.remove(graph_def_file)
def testFrozenGraphDefNonPlaceholder(self):
with ops.Graph().as_default():
in_tensor = random_ops.random_normal(shape=[1, 16, 16, 3], name='random')
_ = in_tensor + in_tensor
sess = session.Session()
# Write graph to file.
graph_def_file = self._getFilepath('model.pb')
write_graph(sess.graph_def, '', graph_def_file, False)
sess.close()
flags_str = ('--graph_def_file={0} --input_arrays={1} '
'--output_arrays={2}'.format(graph_def_file, 'random', 'add'))
self._run(flags_str, should_succeed=True)
os.remove(graph_def_file)
def testQATFrozenGraphDefInt8(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., max=1., name='output',
num_bits=16) # INT8 inference type works for 16 bits fake quant.
sess = session.Session()
# Write graph to file.
graph_def_file = self._getFilepath('model.pb')
write_graph(sess.graph_def, '', graph_def_file, False)
sess.close()
flags_str = ('--inference_type=INT8 --std_dev_values=128,128 '
'--mean_values=128,128 '
'--graph_def_file={0} --input_arrays={1},{2} '
'--output_arrays={3}'.format(graph_def_file, 'inputA',
'inputB', 'output'))
self._run(flags_str, should_succeed=True)
os.remove(graph_def_file)
def testQATFrozenGraphDefUInt8(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., max=1., name='output')
sess = session.Session()
# Write graph to file.
graph_def_file = self._getFilepath('model.pb')
write_graph(sess.graph_def, '', graph_def_file, False)
sess.close()
# Define converter flags
flags_str = ('--std_dev_values=128,128 --mean_values=128,128 '
'--graph_def_file={0} --input_arrays={1} '
'--output_arrays={2}'.format(graph_def_file, 'inputA,inputB',
'output'))
# Set inference_type UINT8 and (default) inference_input_type UINT8
flags_str_1 = flags_str + ' --inference_type=UINT8'
self._run(flags_str_1, should_succeed=True)
# Set inference_type UINT8 and inference_input_type FLOAT
flags_str_2 = flags_str_1 + ' --inference_input_type=FLOAT'
self._run(flags_str_2, should_succeed=True)
os.remove(graph_def_file)
def testSavedModel(self):
saved_model_dir = self._getFilepath('model')
with ops.Graph().as_default():
with session.Session() as sess:
in_tensor = array_ops.placeholder(
shape=[1, 16, 16, 3], dtype=dtypes.float32, name='inputB')
out_tensor = in_tensor + in_tensor
inputs = {'x': in_tensor}
outputs = {'z': out_tensor}
saved_model.simple_save(sess, saved_model_dir, inputs, outputs)
flags_str = '--saved_model_dir={}'.format(saved_model_dir)
self._run(flags_str, should_succeed=True)
def _createSavedModelWithCustomOp(self, opname='CustomAdd'):
custom_opdefs_str = (
'name: \'' + opname + '\' input_arg: {name: \'Input1\' type: DT_FLOAT} '
'input_arg: {name: \'Input2\' type: DT_FLOAT} output_arg: {name: '
'\'Output\' type: DT_FLOAT}')
# Create a graph that has one add op.
new_graph = graph_pb2.GraphDef()
with ops.Graph().as_default():
with session.Session() as sess:
in_tensor = array_ops.placeholder(
shape=[1, 16, 16, 3], dtype=dtypes.float32, name='input')
out_tensor = in_tensor + in_tensor
inputs = {'x': in_tensor}
outputs = {'z': out_tensor}
new_graph.CopyFrom(sess.graph_def)
# Rename Add op name to opname.
for node in new_graph.node:
if node.op.startswith('Add'):
node.op = opname
del node.attr['T']
# Register custom op defs to import modified graph def.
register_custom_opdefs([custom_opdefs_str])
# Store saved model.
saved_model_dir = self._getFilepath('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)
return (saved_model_dir, custom_opdefs_str)
def testEnsureCustomOpdefsFlag(self):
saved_model_dir, _ = self._createSavedModelWithCustomOp()
# Ensure --custom_opdefs.
flags_str = ('--saved_model_dir={0} --allow_custom_ops '
'--experimental_new_converter'.format(saved_model_dir))
self._run(flags_str, should_succeed=False)
def testSavedModelWithCustomOpdefsFlag(self):
saved_model_dir, custom_opdefs_str = self._createSavedModelWithCustomOp()
# Valid conversion.
flags_str = (
'--saved_model_dir={0} --custom_opdefs="{1}" --allow_custom_ops '
'--experimental_new_converter'.format(saved_model_dir,
custom_opdefs_str))
self._run(
flags_str,
should_succeed=True,
expected_ops_in_converted_model=['CustomAdd'])
def testSavedModelWithFlex(self):
saved_model_dir, custom_opdefs_str = self._createSavedModelWithCustomOp(
opname='CustomAdd2')
# Valid conversion. OpDef already registered.
flags_str = ('--saved_model_dir={0} --allow_custom_ops '
'--custom_opdefs="{1}" '
'--experimental_new_converter '
'--experimental_select_user_tf_ops=CustomAdd2 '
'--target_ops=TFLITE_BUILTINS,SELECT_TF_OPS'.format(
saved_model_dir, custom_opdefs_str))
self._run(
flags_str,
should_succeed=True,
expected_ops_in_converted_model=['FlexCustomAdd2'])
def testSavedModelWithInvalidCustomOpdefsFlag(self):
saved_model_dir, _ = self._createSavedModelWithCustomOp()
invalid_custom_opdefs_str = (
'name: \'CustomAdd\' input_arg: {name: \'Input1\' type: DT_FLOAT} '
'output_arg: {name: \'Output\' type: DT_FLOAT}')
# Valid conversion.
flags_str = (
'--saved_model_dir={0} --custom_opdefs="{1}" --allow_custom_ops '
'--experimental_new_converter'.format(saved_model_dir,
invalid_custom_opdefs_str))
self._run(flags_str, should_succeed=False)
def testKerasFile(self):
keras_file = self._getKerasModelFile()
flags_str = '--keras_model_file={}'.format(keras_file)
self._run(flags_str, should_succeed=True)
os.remove(keras_file)
def testKerasFileMLIR(self):
keras_file = self._getKerasModelFile()
flags_str = (
'--keras_model_file={} --experimental_new_converter'.format(keras_file))
self._run(flags_str, should_succeed=True)
os.remove(keras_file)
def _initObjectDetectionArgs(self):
# Initializes the arguments required for the object detection model.
# Looks for the model file which is saved in a different location internally
# and externally.
filename = resource_loader.get_path_to_datafile('testdata/tflite_graph.pb')
if not os.path.exists(filename):
filename = os.path.join(
resource_loader.get_root_dir_with_all_resources(),
'../tflite_mobilenet_ssd_quant_protobuf/tflite_graph.pb')
if not os.path.exists(filename):
raise IOError("File '{0}' does not exist.".format(filename))
self._graph_def_file = filename
self._input_arrays = 'normalized_input_image_tensor'
self._output_arrays = (
'TFLite_Detection_PostProcess,TFLite_Detection_PostProcess:1,'
'TFLite_Detection_PostProcess:2,TFLite_Detection_PostProcess:3')
self._input_shapes = '1,300,300,3'
def testObjectDetection(self):
"""Tests object detection model through TOCO."""
self._initObjectDetectionArgs()
flags_str = ('--graph_def_file={0} --input_arrays={1} '
'--output_arrays={2} --input_shapes={3} '
'--allow_custom_ops'.format(self._graph_def_file,
self._input_arrays,
self._output_arrays,
self._input_shapes))
self._run(flags_str, should_succeed=True)
def testObjectDetectionMLIR(self):
"""Tests object detection model through MLIR converter."""
self._initObjectDetectionArgs()
custom_opdefs_str = (
'name: \'TFLite_Detection_PostProcess\' '
'input_arg: { name: \'raw_outputs/box_encodings\' type: DT_FLOAT } '
'input_arg: { name: \'raw_outputs/class_predictions\' type: DT_FLOAT } '
'input_arg: { name: \'anchors\' type: DT_FLOAT } '
'output_arg: { name: \'TFLite_Detection_PostProcess\' type: DT_FLOAT } '
'output_arg: { name: \'TFLite_Detection_PostProcess:1\' '
'type: DT_FLOAT } '
'output_arg: { name: \'TFLite_Detection_PostProcess:2\' '
'type: DT_FLOAT } '
'output_arg: { name: \'TFLite_Detection_PostProcess:3\' '
'type: DT_FLOAT } '
'attr : { name: \'h_scale\' type: \'float\'} '
'attr : { name: \'max_classes_per_detection\' type: \'int\'} '
'attr : { name: \'max_detections\' type: \'int\'} '
'attr : { name: \'nms_iou_threshold\' type: \'float\'} '
'attr : { name: \'nms_score_threshold\' type: \'float\'} '
'attr : { name: \'num_classes\' type: \'int\'} '
'attr : { name: \'w_scale\' type: \'float\'} '
'attr : { name: \'x_scale\' type: \'float\'} '
'attr : { name: \'y_scale\' type: \'float\'}')
flags_str = ('--graph_def_file={0} --input_arrays={1} '
'--output_arrays={2} --input_shapes={3} '
'--custom_opdefs="{4}"'.format(self._graph_def_file,
self._input_arrays,
self._output_arrays,
self._input_shapes,
custom_opdefs_str))
# Ensure --allow_custom_ops.
flags_str_final = ('{} --allow_custom_ops').format(flags_str)
self._run(
flags_str_final,
should_succeed=True,
expected_ops_in_converted_model=['TFLite_Detection_PostProcess'])
def testObjectDetectionMLIRWithFlex(self):
"""Tests object detection model through MLIR converter."""
self._initObjectDetectionArgs()
flags_str = ('--graph_def_file={0} --input_arrays={1} '
'--output_arrays={2} --input_shapes={3}'.format(
self._graph_def_file, self._input_arrays,
self._output_arrays, self._input_shapes))
# Valid conversion.
flags_str_final = (
'{} --allow_custom_ops '
'--experimental_new_converter '
'--experimental_select_user_tf_ops=TFLite_Detection_PostProcess '
'--target_ops=TFLITE_BUILTINS,SELECT_TF_OPS').format(flags_str)
self._run(
flags_str_final,
should_succeed=True,
expected_ops_in_converted_model=['FlexTFLite_Detection_PostProcess'])
class TfLiteConvertV2Test(TestModels):
@test_util.run_v2_only
def testSavedModel(self):
input_data = constant_op.constant(1., shape=[1])
root = autotrackable.AutoTrackable()
root.f = def_function.function(lambda x: 2. * x)
to_save = root.f.get_concrete_function(input_data)
saved_model_dir = self._getFilepath('model')
save(root, saved_model_dir, to_save)
flags_str = '--saved_model_dir={}'.format(saved_model_dir)
self._run(flags_str, should_succeed=True)
@test_util.run_v2_only
def testKerasFile(self):
keras_file = self._getKerasModelFile()
flags_str = '--keras_model_file={}'.format(keras_file)
self._run(flags_str, should_succeed=True)
os.remove(keras_file)
@test_util.run_v2_only
def testKerasFileMLIR(self):
keras_file = self._getKerasModelFile()
flags_str = (
'--keras_model_file={} --experimental_new_converter'.format(keras_file))
self._run(flags_str, should_succeed=True)
os.remove(keras_file)
@test_util.run_v2_only
def testFunctionalKerasModel(self):
keras_file = self._getKerasFunctionalModelFile()
flags_str = '--keras_model_file={}'.format(keras_file)
self._run(flags_str, should_succeed=True,
expected_output_shapes=[[1, 1], [1, 2]])
os.remove(keras_file)
@test_util.run_v2_only
def testFunctionalKerasModelMLIR(self):
keras_file = self._getKerasFunctionalModelFile()
flags_str = (
'--keras_model_file={} --experimental_new_converter'.format(keras_file))
self._run(flags_str, should_succeed=True,
expected_output_shapes=[[1, 1], [1, 2]])
os.remove(keras_file)
def testMissingRequired(self):
self._run('--invalid_args', should_succeed=False)
def testMutuallyExclusive(self):
self._run(
'--keras_model_file=model.h5 --saved_model_dir=/tmp/',
should_succeed=False)
class ArgParserTest(test_util.TensorFlowTestCase, parameterized.TestCase):
@parameterized.named_parameters(('v1', False), ('v2', True))
def test_without_experimental_new_converter(self, use_v2_converter):
args = [
'--saved_model_dir=/tmp/saved_model/',
'--output_file=/tmp/output.tflite',
]
# Note that when the flag parses to None, the converter uses the default
# value, which is True.
parser = tflite_convert._get_parser(use_v2_converter=use_v2_converter)
parsed_args = parser.parse_args(args)
self.assertTrue(parsed_args.experimental_new_converter)
self.assertIsNone(parsed_args.experimental_new_quantizer)
@parameterized.named_parameters(('v1', False), ('v2', True))
def test_experimental_new_converter_none(self, use_v2_converter):
args = [
'--saved_model_dir=/tmp/saved_model/',
'--output_file=/tmp/output.tflite',
'--experimental_new_converter',
]
parser = tflite_convert._get_parser(use_v2_converter=use_v2_converter)
parsed_args = parser.parse_args(args)
self.assertTrue(parsed_args.experimental_new_converter)
@parameterized.named_parameters(
('v1_true', False, True),
('v1_false', False, False),
('v2_true', True, True),
('v2_false', True, False),
)
def test_experimental_new_converter(self, use_v2_converter, new_converter):
args = [
'--saved_model_dir=/tmp/saved_model/',
'--output_file=/tmp/output.tflite',
'--experimental_new_converter={}'.format(new_converter),
]
parser = tflite_convert._get_parser(use_v2_converter=use_v2_converter)
parsed_args = parser.parse_args(args)
self.assertEqual(parsed_args.experimental_new_converter, new_converter)
@parameterized.named_parameters(('v1', False), ('v2', True))
def test_experimental_new_quantizer_none(self, use_v2_converter):
args = [
'--saved_model_dir=/tmp/saved_model/',
'--output_file=/tmp/output.tflite',
'--experimental_new_quantizer',
]
parser = tflite_convert._get_parser(use_v2_converter=use_v2_converter)
parsed_args = parser.parse_args(args)
self.assertTrue(parsed_args.experimental_new_quantizer)
@parameterized.named_parameters(
('v1_true', False, True),
('v1_false', False, False),
('v2_true', True, True),
('v2_false', True, False),
)
def test_experimental_new_quantizer(self, use_v2_converter, new_quantizer):
args = [
'--saved_model_dir=/tmp/saved_model/',
'--output_file=/tmp/output.tflite',
'--experimental_new_quantizer={}'.format(new_quantizer),
]
parser = tflite_convert._get_parser(use_v2_converter=use_v2_converter)
parsed_args = parser.parse_args(args)
self.assertEqual(parsed_args.experimental_new_quantizer, new_quantizer)
if __name__ == '__main__':
test.main()
tensorflow/lite/python/tflite_keras_util.py
+234
View File
@@ -0,0 +1,234 @@
# Copyright 2020 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.
# ==============================================================================
"""Keras functions required by TensorFlow Lite.
The functions defined in this library have been copied over from Keras in order
to remove the dependency from TensorFlow Lite to Keras. The functions which
could not be copied over are accessed using the dependency inversion principle.
(for details, refer to tensorflow/python/util/keras_deps.py).
"""
import copy
from tensorflow.python.eager import def_function
from tensorflow.python.framework import tensor_spec
from tensorflow.python.util import keras_deps
from tensorflow.python.util import nest
from tensorflow.python.util.compat import collections_abc
def _enforce_names_consistency(specs):
"""Enforces that either all specs have names or none do."""
def _has_name(spec):
return hasattr(spec, 'name') and spec.name is not None
def _clear_name(spec):
spec = copy.deepcopy(spec)
if hasattr(spec, 'name'):
spec._name = None # pylint:disable=protected-access
return spec
flat_specs = nest.flatten(specs)
name_inconsistency = (
any(_has_name(s) for s in flat_specs) and
not all(_has_name(s) for s in flat_specs))
if name_inconsistency:
specs = nest.map_structure(_clear_name, specs)
return specs
def get_save_spec(model):
"""Returns the save spec of the subclassing keras model."""
shapes_dict = getattr(model, '_build_shapes_dict', None)
if not shapes_dict:
return None
if 'input_shape' not in shapes_dict:
raise ValueError(
'Model {} cannot be saved because the input shapes have not been set.'
)
input_shape = shapes_dict['input_shape']
if isinstance(input_shape, tuple):
shape = input_shape
shape = (None,) + shape[1:]
return tensor_spec.TensorSpec(
shape=shape, dtype=model.input_dtype
)
elif isinstance(input_shape, dict):
specs = {}
for key, shape in input_shape.items():
shape = (None,) + shape[1:]
specs[key] = tensor_spec.TensorSpec(
shape=shape, dtype=model.input_dtype, name=key
)
return specs
elif isinstance(input_shape, list):
specs = []
for shape in input_shape:
shape = (None,) + shape[1:]
specs.append(tensor_spec.TensorSpec(shape=shape, dtype=model.input_dtype))
return specs
def model_input_signature(model, keep_original_batch_size=False):
"""Inspect model to get its input signature.
The model's input signature is a list with a single (possibly-nested) object.
This is due to the Keras-enforced restriction that tensor inputs must be
passed in as the first argument.
For example, a model with input {'feature1': <Tensor>, 'feature2': <Tensor>}
will have input signature: [{'feature1': TensorSpec, 'feature2': TensorSpec}]
Args:
model: Keras Model object.
keep_original_batch_size: A boolean indicating whether we want to keep using
the original batch size or set it to None. Default is `False`, which means
that the batch dim of the returned input signature will always be set to
`None`.
Returns:
A list containing either a single TensorSpec or an object with nested
TensorSpecs. This list does not contain the `training` argument.
"""
if hasattr(model, 'save_spec'):
input_specs = model.save_spec(dynamic_batch=not keep_original_batch_size)
if input_specs is None:
return None
# The model's save spec returns (args, kwargs). Extract the first input arg
# to use as the input spec.
# TODO(b/188105669): Add support for multiple tensor arguments.
input_specs = input_specs[0][0]
else:
input_specs = model._get_save_spec( # pylint: disable=protected-access
dynamic_batch=not keep_original_batch_size)
if input_specs is None:
return None
input_specs = _enforce_names_consistency(input_specs)
# Return a list with a single element as the model's input signature.
if isinstance(input_specs,
collections_abc.Sequence) and len(input_specs) == 1:
# Note that the isinstance check filters out single-element dictionaries,
# which should also be wrapped as a single-element list.
return input_specs
else:
return [input_specs]
def raise_model_input_error(model):
raise ValueError(
'Model {} cannot be saved because the input shapes have not been '
'set. Usually, input shapes are automatically determined from calling'
' `.fit()` or `.predict()`. To manually set the shapes, call '
'`model.build(input_shape)`.'.format(model))
def _create_pseudo_names(tensors, prefix):
"""Creates pseudo {input | output} names for subclassed Models.
Warning: this function should only be used to define default
names for `Metics` and `SavedModel`. No other use cases should
rely on a `Model`'s input or output names.
Example with dict:
`{'a': [x1, x2], 'b': x3}` becomes:
`['a_1', 'a_2', 'b']`
Example with list:
`[x, y]` becomes:
`['output_1', 'output_2']`
Args:
tensors: `Model`'s outputs or inputs.
prefix: 'output_' for outputs, 'input_' for inputs.
Returns:
Flattened list of pseudo names.
"""
def one_index(ele):
# Start with "output_1" instead of "output_0".
if isinstance(ele, int):
return ele + 1
return ele
flat_paths = list(nest.yield_flat_paths(tensors))
flat_paths = nest.map_structure(one_index, flat_paths)
names = []
for path in flat_paths:
if not path:
name = prefix + '1' # Single output.
else:
name = '_'.join(str(p) for p in path)
if isinstance(path[0], int):
name = prefix + name
names.append(name)
return names
def create_pseudo_output_names(outputs):
"""Create pseudo output names for a subclassed Model."""
return _create_pseudo_names(outputs, prefix='output_')
def trace_model_call(model, input_signature=None):
"""Trace the model call to create a tf.function for exporting a Keras model.
Args:
model: A Keras model.
input_signature: optional, a list of tf.TensorSpec objects specifying the
inputs to the model.
Returns:
A tf.function wrapping the model's call function with input signatures set.
Raises:
ValueError: if input signature cannot be inferred from the model.
"""
if input_signature is None:
if isinstance(model.call, def_function.Function):
input_signature = model.call.input_signature
if input_signature is None:
input_signature = model_input_signature(model)
if input_signature is None:
raise_model_input_error(model)
@def_function.function(input_signature=input_signature, autograph=False)
def _wrapped_model(*args):
"""A concrete tf.function that wraps the model's call function."""
# When given a single input, Keras models will call the model on the tensor
# rather than a list consisting of the single tensor.
inputs = args[0] if len(input_signature) == 1 else list(args)
with keras_deps.get_call_context_function()().enter(
model,
inputs=inputs,
build_graph=False,
call_context_args={'training': False},
saving=True,
):
outputs = model(inputs, training=False)
return outputs
return _wrapped_model
tensorflow/lite/python/util.py
+1177
View File
File diff suppressed because it is too large Load Diff
tensorflow/lite/python/util_test.py
+518
View File
@@ -0,0 +1,518 @@
# 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.
# ==============================================================================
"""Tests for util.py."""
import os
from absl.testing import parameterized
import numpy as np
import tensorflow as tf
from tensorflow.lite.python import lite
from tensorflow.lite.python import util
from tensorflow.lite.tools.flatbuffer_utils import read_model as _read_model
from tensorflow.python.client import session
from tensorflow.python.framework import convert_to_constants
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import test_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import while_loop
from tensorflow.python.platform import test
# TODO(nupurgarg): Add test for Grappler and frozen graph related functions.
class UtilTest(test_util.TensorFlowTestCase):
def testConvertEnumToDtype(self):
self.assertEqual(
util._convert_tflite_enum_type_to_tf_type(0), dtypes.float32)
self.assertEqual(
util._convert_tflite_enum_type_to_tf_type(1), dtypes.float16)
self.assertEqual(util._convert_tflite_enum_type_to_tf_type(2), dtypes.int32)
self.assertEqual(util._convert_tflite_enum_type_to_tf_type(3), dtypes.uint8)
self.assertEqual(util._convert_tflite_enum_type_to_tf_type(4), dtypes.int64)
self.assertEqual(
util._convert_tflite_enum_type_to_tf_type(5), dtypes.string)
self.assertEqual(util._convert_tflite_enum_type_to_tf_type(6), dtypes.bool)
self.assertEqual(util._convert_tflite_enum_type_to_tf_type(7), dtypes.int16)
self.assertEqual(
util._convert_tflite_enum_type_to_tf_type(8), dtypes.complex64)
self.assertEqual(util._convert_tflite_enum_type_to_tf_type(9), dtypes.int8)
self.assertEqual(
util._convert_tflite_enum_type_to_tf_type(10), dtypes.float64)
self.assertEqual(
util._convert_tflite_enum_type_to_tf_type(11), dtypes.complex128)
self.assertEqual(
util._convert_tflite_enum_type_to_tf_type(16), dtypes.uint32)
with self.assertRaises(ValueError) as error:
util._convert_tflite_enum_type_to_tf_type(20)
self.assertEqual(
"Unsupported enum 20. The valid map of enum to tf types is : "
"{0: tf.float32, 1: tf.float16, 2: tf.int32, 3: tf.uint8, 4: tf.int64, "
"5: tf.string, 6: tf.bool, 7: tf.int16, 8: tf.complex64, 9: tf.int8, "
"10: tf.float64, 11: tf.complex128, 16: tf.uint32}",
str(error.exception))
def testTensorName(self):
with ops.Graph().as_default():
in_tensor = array_ops.placeholder(dtype=dtypes.float32, shape=[4])
out_tensors = array_ops.split(
value=in_tensor, num_or_size_splits=[1, 1, 1, 1], axis=0)
expect_names = ["split", "split:1", "split:2", "split:3"]
for i in range(len(expect_names)):
got_name = util.get_tensor_name(out_tensors[i])
self.assertEqual(got_name, expect_names[i])
def testUint32PassThrough(self):
model = tf.keras.Sequential([
tf.keras.layers.InputLayer(input_shape=(4,), dtype=tf.uint32),
tf.keras.layers.Reshape(target_shape=(2, 2))
])
converter = lite.TFLiteConverterV2.from_keras_model(model)
tflite_model = converter.convert()
interpreter = lite.Interpreter(model_content=tflite_model)
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()[0]
output_details = interpreter.get_output_details()[0]
self.assertEqual(input_details["dtype"], np.uint32)
self.assertEqual(output_details["dtype"], np.uint32)
in_array = np.array([[1, 1, 1, 1]], dtype="uint32") * ((1 << 32) - 1)
expected_out = np.reshape(in_array, (2, 2))
interpreter.set_tensor(input_details["index"], in_array)
interpreter.invoke()
output_data = interpreter.get_tensor(output_details["index"])[0]
self.assertAllEqual(expected_out, output_data)
@test_util.enable_control_flow_v2
def testRemoveLowerUsingSwitchMerge(self):
with ops.Graph().as_default():
i = array_ops.placeholder(dtype=dtypes.int32, shape=())
c = lambda i: math_ops.less(i, 10)
b = lambda i: math_ops.add(i, 1)
while_loop.while_loop(c, b, [i])
sess = session.Session()
new_graph_def = convert_to_constants.disable_lower_using_switch_merge(
sess.graph_def)
lower_using_switch_merge_is_removed = False
for node in new_graph_def.node:
if node.op == "While" or node.op == "StatelessWhile":
if not node.attr["_lower_using_switch_merge"].b:
lower_using_switch_merge_is_removed = True
self.assertTrue(lower_using_switch_merge_is_removed)
def testConvertBytes(self):
source, header = util.convert_bytes_to_c_source(
b"\x00\x01\x02\x23", "foo", 16, use_tensorflow_license=False)
self.assertTrue(
source.find("const unsigned char foo[] DATA_ALIGN_ATTRIBUTE = {"))
self.assertTrue(source.find(""" 0x00, 0x01,
0x02, 0x23,"""))
self.assertNotEqual(-1, source.find("const int foo_len = 4;"))
self.assertEqual(-1, source.find("/* Copyright"))
self.assertEqual(-1, source.find("#include " ""))
self.assertNotEqual(-1, header.find("extern const unsigned char foo[];"))
self.assertNotEqual(-1, header.find("extern const int foo_len;"))
self.assertEqual(-1, header.find("/* Copyright"))
source, header = util.convert_bytes_to_c_source(
b"\xff\xfe\xfd\xfc",
"bar",
80,
include_guard="MY_GUARD",
include_path="my/guard.h",
use_tensorflow_license=True)
self.assertNotEqual(
-1, source.find("const unsigned char bar[] DATA_ALIGN_ATTRIBUTE = {"))
self.assertNotEqual(-1, source.find(""" 0xff, 0xfe, 0xfd, 0xfc,"""))
self.assertNotEqual(-1, source.find("/* Copyright"))
self.assertNotEqual(-1, source.find("#include \"my/guard.h\""))
self.assertNotEqual(-1, header.find("#ifndef MY_GUARD"))
self.assertNotEqual(-1, header.find("#define MY_GUARD"))
self.assertNotEqual(-1, header.find("/* Copyright"))
class TensorFunctionsTest(test_util.TensorFlowTestCase):
def testGetTensorsValid(self):
with ops.Graph().as_default():
in_tensor = array_ops.placeholder(
dtype=dtypes.float32, shape=[1, 16, 16, 3])
_ = in_tensor + in_tensor
sess = session.Session()
tensors = util.get_tensors_from_tensor_names(sess.graph, ["Placeholder"])
self.assertEqual("Placeholder:0", tensors[0].name)
def testGetTensorsInvalid(self):
with ops.Graph().as_default():
in_tensor = array_ops.placeholder(
dtype=dtypes.float32, shape=[1, 16, 16, 3])
_ = in_tensor + in_tensor
sess = session.Session()
with self.assertRaises(ValueError) as error:
util.get_tensors_from_tensor_names(sess.graph, ["invalid-input"])
self.assertEqual("Invalid tensors 'invalid-input' were found.",
str(error.exception))
def testSetTensorShapeValid(self):
with ops.Graph().as_default():
tensor = array_ops.placeholder(dtype=dtypes.float32, shape=[None, 3, 5])
self.assertAllEqual([None, 3, 5], tensor.shape)
util.set_tensor_shapes([tensor], {"Placeholder": [5, 3, 5]})
self.assertAllEqual([5, 3, 5], tensor.shape)
def testSetTensorShapeNoneValid(self):
with ops.Graph().as_default():
tensor = array_ops.placeholder(dtype=dtypes.float32)
util.set_tensor_shapes([tensor], {"Placeholder": [1, 3, 5]})
self.assertAllEqual([1, 3, 5], tensor.shape)
def testSetTensorShapeArrayInvalid(self):
# Tests set_tensor_shape where the tensor name passed in doesn't exist.
with ops.Graph().as_default():
tensor = array_ops.placeholder(dtype=dtypes.float32, shape=[None, 3, 5])
self.assertAllEqual([None, 3, 5], tensor.shape)
with self.assertRaises(ValueError) as error:
util.set_tensor_shapes([tensor], {"invalid-input": [5, 3, 5]})
self.assertEqual(
"Invalid tensor 'invalid-input' found in tensor shapes map.",
str(error.exception))
self.assertAllEqual([None, 3, 5], tensor.shape)
def testSetTensorShapeDimensionInvalid(self):
# Tests set_tensor_shape where the shape passed in is incompatible.
with ops.Graph().as_default():
tensor = array_ops.placeholder(dtype=dtypes.float32, shape=[None, 3, 5])
self.assertAllEqual([None, 3, 5], tensor.shape)
with self.assertRaises(ValueError) as error:
util.set_tensor_shapes([tensor], {"Placeholder": [1, 5, 5]})
self.assertIn("The shape of tensor 'Placeholder' cannot be changed",
str(error.exception))
self.assertAllEqual([None, 3, 5], tensor.shape)
def testSetTensorShapeEmpty(self):
with ops.Graph().as_default():
tensor = array_ops.placeholder(dtype=dtypes.float32, shape=[None, 3, 5])
self.assertAllEqual([None, 3, 5], tensor.shape)
util.set_tensor_shapes([tensor], {})
self.assertAllEqual([None, 3, 5], tensor.shape)
def _get_keras_model(add_unquantizable_layer=False):
"""Define Sample keras model and returns it."""
# Define a pseudo MNIST dataset (as downloading the dataset on-the-fly causes
# network connection failures)
n = 10 # Number of samples
images = np.random.randint(low=0, high=255, size=[n, 28, 28], dtype=np.uint8)
labels = np.random.randint(low=0, high=9, size=(n,), dtype=np.uint8)
# Normalize the input image so that each pixel value is between 0 to 1.
images = images / 255.0
# Define TF model
model = tf.keras.Sequential([
tf.keras.layers.InputLayer(input_shape=(28, 28)),
tf.keras.layers.Reshape(target_shape=(28, 28, 1)),
tf.keras.layers.Conv2D(filters=12, kernel_size=(3, 3), activation="relu"),
tf.keras.layers.MaxPooling2D(pool_size=(2, 2)),
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(10)
])
if add_unquantizable_layer:
# This adds Neg op to the model which will remain as float.
model.add(tf.keras.layers.Lambda(lambda x: -x))
# Train
model.compile(
optimizer="adam",
loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
metrics=["accuracy"])
model.fit(
images,
labels,
epochs=1,
validation_split=0.1,
)
return model
def _generate_integer_tflite_model(quantization_type=dtypes.int8,
use_saved_model=False,
saved_model_dir=None,
add_unquantizable_layer=False):
"""Define an integer post-training quantized tflite model."""
model = _get_keras_model(add_unquantizable_layer)
if not use_saved_model:
# Convert TF Model to an Integer Quantized TFLite Model
converter = lite.TFLiteConverterV2.from_keras_model(model)
else:
tf.saved_model.save(model, saved_model_dir)
converter = lite.TFLiteConverterV2.from_saved_model(saved_model_dir)
converter.optimizations = {lite.Optimize.DEFAULT}
def representative_dataset_gen():
for _ in range(2):
yield [
np.random.uniform(low=0, high=1, size=(1, 28, 28)).astype(np.float32)
]
converter.representative_dataset = representative_dataset_gen
if quantization_type == dtypes.int8:
converter.target_spec.supported_ops = {lite.OpsSet.TFLITE_BUILTINS_INT8}
else:
converter.target_spec.supported_ops = {
lite.OpsSet
.EXPERIMENTAL_TFLITE_BUILTINS_ACTIVATIONS_INT16_WEIGHTS_INT8
}
tflite_model = converter.convert()
return tflite_model
def _test_param_modify_integer_model_io_type():
"""Function to generate parameterized inputs for testing."""
params = []
str_template = "_{}{}{}{}"
map_model_type = {
"PostTraining": True,
# "DuringTraining": False,
}
map_quantize_type_to_io_types = {
tf.int8: {tf.float32, tf.int8, tf.uint8},
tf.int16: {tf.float32, tf.int16}
}
for k1, v1 in map_model_type.items():
for qtype, v2 in map_quantize_type_to_io_types.items():
qstr = "_IntegerQuantize{}".format(qtype.name.capitalize())
for itype in v2:
istr = "_Input{}".format(itype.name.capitalize())
for otype in v2:
ostr = "_Output{}".format(otype.name.capitalize())
params.append((str_template.format(k1, qstr, istr,
ostr), v1, qtype, itype, otype))
return params
class UtilModifyIntegerQuantizedModelIOTypeTest(test_util.TensorFlowTestCase,
parameterized.TestCase):
@classmethod
def setUpClass(cls):
super(UtilModifyIntegerQuantizedModelIOTypeTest, cls).setUpClass()
cls.post_train_int8_model = _generate_integer_tflite_model()
cls.post_train_int16_model = _generate_integer_tflite_model(
quantization_type=dtypes.int16)
@parameterized.named_parameters(_test_param_modify_integer_model_io_type())
def test(self, is_post_train, quantization_type, in_tftype, out_tftype):
"""Modify the float input/output type of an integer quantized model."""
def _run_tflite_inference(model, in_tftype, out_tftype):
"""Run inference on a model with a specific input/output type."""
# Load TFLite model and allocate tensors.
interpreter = lite.Interpreter(model_content=model)
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()[0]
output_details = interpreter.get_output_details()[0]
# Validate TFLite model input and output types
self.assertEqual(input_details["dtype"], in_tftype.as_numpy_dtype)
self.assertEqual(output_details["dtype"], out_tftype.as_numpy_dtype)
# Define Input
np.random.seed(0)
input_data = np.random.uniform(low=0, high=1, size=(1, 28, 28))
input_data = input_data.astype(np.float32)
if input_details["dtype"] != np.float32:
# quantize float to int
scale, zero_point = input_details["quantization"]
input_data = input_data / scale + zero_point
input_data = input_data.astype(input_details["dtype"])
# Run Inference
interpreter.set_tensor(input_details["index"], input_data)
interpreter.invoke()
# Get output
output_data = interpreter.get_tensor(output_details["index"])[0]
if output_details["dtype"] != np.float32:
# dequantize int to float
scale, zero_point = output_details["quantization"]
output_data = output_data.astype(np.float32)
output_data = (output_data - zero_point) * scale
return output_data
if is_post_train and quantization_type == tf.int8:
model = self.__class__.post_train_int8_model
elif is_post_train and quantization_type == tf.int16:
model = self.__class__.post_train_int16_model
else:
model = None
# Run model inference with float input output type
output_data = _run_tflite_inference(model, tf.float32, tf.float32)
# Modify the model io types to the target input/output types.
model_io = util.modify_model_io_type(model, in_tftype, out_tftype)
# Run model inference with modified integer input output type
output_io_data = _run_tflite_inference(model_io, in_tftype, out_tftype)
# Validate that both the outputs are the same
self.assertAllClose(output_data, output_io_data, atol=1.0)
# Modify the model with the target input/output types should be a no op.
model_io = util.modify_model_io_type(model_io, in_tftype, out_tftype)
# Run model inference with modified integer input output type
output_io_data = _run_tflite_inference(model_io, in_tftype, out_tftype)
# Validate that both the outputs are the same
self.assertAllClose(output_data, output_io_data, atol=1.0)
class UtilModifyIntegerQuantizedModelIOTypeSignatureDefTest(
test_util.TensorFlowTestCase):
def _generate_integer_tflite_model_from_saved_model(self):
"""Define an integer post-training quantized model from saved model."""
saved_model_dir = os.path.join(self.get_temp_dir(), "simple_savedmodel")
return _generate_integer_tflite_model(
use_saved_model=True,
saved_model_dir=saved_model_dir,
add_unquantizable_layer=True)
def test(self):
"""Makes sure modifying IO types updates Signature correctly."""
post_train_int8_model = (
self._generate_integer_tflite_model_from_saved_model())
modified_model = util.modify_model_io_type(post_train_int8_model, tf.int8,
tf.float32)
interpreter = lite.Interpreter(model_content=modified_model)
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
signature = interpreter._get_full_signature_list()
input_ids = []
output_ids = []
for input_tensor in input_details:
input_ids.append(input_tensor["index"])
for output_tensor in output_details:
output_ids.append(output_tensor["index"])
for _, tensor_id in signature["serving_default"]["inputs"].items():
assert tensor_id in input_ids
for _, tensor_id in signature["serving_default"]["outputs"].items():
assert tensor_id in output_ids
class UtilModifyIntegerQuantizedConcatResidualModelIOTypeTest(
test_util.TensorFlowTestCase, parameterized.TestCase
):
def _generate_int8_f32io_concat_residual_tflite(self, number_of_inputs=3):
dtype = float
class ConcatNResidual(tf.keras.layers.Layer):
"""A simple concat and residual Keras Model."""
def __init__(self, number_of_inputs=3, **kwargs):
super().__init__(**kwargs)
self.number_of_inputs = number_of_inputs
self.conv = tf.keras.layers.Conv2D(2, (2, 2), padding="same")
self.mins = [-0.01 * (i + 1) for i in range(self.number_of_inputs)]
self.maxs = [0.01 * (i + 1) for i in range(self.number_of_inputs)]
def call(self, inputs):
xs = [
tf.quantization.fake_quant_with_min_max_args(
inputs[i], self.mins[i], self.maxs[i]
)
for i in range(self.number_of_inputs)
]
x = tf.keras.backend.concatenate(xs, 1)
x = x[:, : inputs[-1].shape[1]]
x = x + xs[-1]
x = tf.quantization.fake_quant_with_min_max_args(x, -2.242, 2.242)
return x
inputs = [
tf.keras.layers.Input(shape=(2, 2, 2), batch_size=1, dtype=dtype)
for _ in range(number_of_inputs)
]
outputs = ConcatNResidual(number_of_inputs)(inputs)
model = tf.keras.Model(inputs, outputs)
converter = lite.TFLiteConverterV2.from_keras_model(model)
converter.optimizations = [lite.Optimize.DEFAULT]
tflite_model = converter.convert()
return tflite_model
def _verify_tensor_connections(self, flatbuffer_model):
"""Verify that all the tensors have input and output ops except the tensors have buffer data."""
tflite_subgraph = flatbuffer_model.subgraphs[0]
tensors = tflite_subgraph.tensors
buffers = flatbuffer_model.buffers
tensors_used_as_inputs = set()
tensors_used_as_outputs = set()
for op in tflite_subgraph.operators:
tensors_used_as_inputs.update(
idx for idx in op.inputs if buffers[tensors[idx].buffer].data is None
)
tensors_used_as_outputs.update(idx for idx in op.outputs)
tensors_used_as_inputs.update(idx for idx in tflite_subgraph.outputs)
tensors_used_as_outputs.update(idx for idx in tflite_subgraph.inputs)
self.assertEqual(tensors_used_as_inputs, tensors_used_as_outputs)
@parameterized.named_parameters([
("_IntOnly_Float32InputOutput", tf.float32),
("_IntOnly_INT8InputOutput", tf.int8),
("_IntOnly_UINT8InputOutput", tf.uint8),
])
def test(self, inference_input_output_type):
"""Make sure modifying IO types removes tensors correctly."""
srqed_int8_f32io_model = self._generate_int8_f32io_concat_residual_tflite()
if inference_input_output_type != tf.float32:
target_model = util.modify_model_io_type(
srqed_int8_f32io_model,
inference_input_output_type,
inference_input_output_type,
)
else:
target_model = srqed_int8_f32io_model
tflite_path = os.path.join(self.get_temp_dir(), "concat_residual.tflite")
with tf.io.gfile.GFile(tflite_path, "wb") as writer:
writer.write(target_model)
flatbuffer_model = _read_model(tflite_path)
self._verify_tensor_connections(flatbuffer_model)
if __name__ == "__main__":
test.main()
tensorflow/lite/python/wrap_toco.py
+33
View File
@@ -0,0 +1,33 @@
# 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.
# ==============================================================================
"""Wraps toco interface with python lazy loader."""
# We need to import pywrap_tensorflow prior to the toco wrapper.
# pylint: disable=invalid-import-order,g-bad-import-order
from tensorflow.python import pywrap_tensorflow # pylint: disable=unused-import
from tensorflow.python import _pywrap_toco_api
def wrapped_toco_convert(
model_flags_str,
toco_flags_str,
input_data_str,
):
"""Wraps TocoConvert with lazy loader."""
return _pywrap_toco_api.TocoConvert(
model_flags_str,
toco_flags_str,
input_data_str,
False, # extended_return
)