421 lines
12 KiB
Python
421 lines
12 KiB
Python
# Description:
|
|
# Tools for manipulating TensorFlow graphs.
|
|
|
|
load("//tensorflow:tensorflow.bzl", "if_xla_available", "py_binary", "py_test", "tf_cc_test")
|
|
load("//tensorflow/python/tools:tools.bzl", "saved_model_compile_aot")
|
|
|
|
package(
|
|
default_visibility = ["//visibility:public"],
|
|
licenses = ["notice"], # Apache 2.0
|
|
)
|
|
|
|
exports_files(["LICENSE"])
|
|
|
|
# Transitive dependencies of this target will be included in the pip package.
|
|
py_library(
|
|
name = "tools_pip",
|
|
data = [
|
|
":freeze_graph",
|
|
":import_pb_to_tensorboard",
|
|
":inspect_checkpoint",
|
|
":optimize_for_inference",
|
|
":print_selective_registration_header",
|
|
":saved_model_cli",
|
|
":strip_unused",
|
|
# Include the TF upgrade script to users can run it directly after install TF
|
|
"//tensorflow/tools/compatibility:tf_upgrade_v2",
|
|
],
|
|
deps = [
|
|
":saved_model_aot_compile",
|
|
":saved_model_utils",
|
|
# The following py_library are needed because
|
|
# py_binary may not depend on them when --define=no_tensorflow_py_deps=true
|
|
# is specified. See https://github.com/tensorflow/tensorflow/issues/22390
|
|
":freeze_graph_lib",
|
|
":optimize_for_inference_lib",
|
|
":selective_registration_header_lib",
|
|
":strip_unused_lib",
|
|
],
|
|
)
|
|
|
|
py_library(
|
|
name = "saved_model_utils",
|
|
srcs = ["saved_model_utils.py"],
|
|
srcs_version = "PY2AND3",
|
|
)
|
|
|
|
py_test(
|
|
name = "saved_model_utils_test",
|
|
size = "small",
|
|
srcs = ["saved_model_utils_test.py"],
|
|
python_version = "PY3",
|
|
srcs_version = "PY2AND3",
|
|
tags = ["no_windows"], # TODO: needs investigation on Windows
|
|
visibility = ["//visibility:private"],
|
|
deps = [
|
|
":saved_model_utils",
|
|
"//tensorflow/python:client_testlib",
|
|
"//tensorflow/python/saved_model",
|
|
],
|
|
)
|
|
|
|
py_library(
|
|
name = "freeze_graph_lib",
|
|
srcs = ["freeze_graph.py"],
|
|
srcs_version = "PY2AND3",
|
|
deps = [
|
|
":saved_model_utils",
|
|
"//tensorflow/core:protos_all_py",
|
|
"//tensorflow/python:client",
|
|
"//tensorflow/python:framework",
|
|
"//tensorflow/python:no_contrib", # TODO(b/34059704): remove when fixed
|
|
"//tensorflow/python:parsing_ops",
|
|
"//tensorflow/python:platform",
|
|
"//tensorflow/python:py_checkpoint_reader",
|
|
"//tensorflow/python:training",
|
|
"//tensorflow/python/estimator:estimator_py",
|
|
"//tensorflow/python/saved_model:loader",
|
|
"@six_archive//:six",
|
|
],
|
|
)
|
|
|
|
py_binary(
|
|
name = "freeze_graph",
|
|
srcs = ["freeze_graph.py"],
|
|
python_version = "PY3",
|
|
srcs_version = "PY2AND3",
|
|
deps = [":freeze_graph_main_lib"],
|
|
)
|
|
|
|
py_library(
|
|
name = "freeze_graph_main_lib",
|
|
srcs = ["freeze_graph.py"],
|
|
srcs_version = "PY2AND3",
|
|
deps = [
|
|
":freeze_graph_lib",
|
|
],
|
|
)
|
|
|
|
py_binary(
|
|
name = "import_pb_to_tensorboard",
|
|
srcs = ["import_pb_to_tensorboard.py"],
|
|
python_version = "PY3",
|
|
srcs_version = "PY2AND3",
|
|
deps = [":import_pb_to_tensorboard_lib"],
|
|
)
|
|
|
|
py_library(
|
|
name = "import_pb_to_tensorboard_lib",
|
|
srcs = ["import_pb_to_tensorboard.py"],
|
|
srcs_version = "PY2AND3",
|
|
deps = [
|
|
"//tensorflow/core:protos_all_py",
|
|
"//tensorflow/python",
|
|
"//tensorflow/python:client",
|
|
"//tensorflow/python:framework",
|
|
"//tensorflow/python:framework_ops",
|
|
"//tensorflow/python:platform",
|
|
"//tensorflow/python:summary",
|
|
],
|
|
)
|
|
|
|
py_test(
|
|
name = "freeze_graph_test",
|
|
size = "small",
|
|
srcs = ["freeze_graph_test.py"],
|
|
python_version = "PY3",
|
|
srcs_version = "PY2AND3",
|
|
deps = [
|
|
":freeze_graph_lib",
|
|
"//tensorflow/core:protos_all_py",
|
|
"//tensorflow/python:client",
|
|
"//tensorflow/python:client_testlib",
|
|
"//tensorflow/python:framework",
|
|
"//tensorflow/python:framework_for_generated_wrappers",
|
|
"//tensorflow/python:framework_test_lib",
|
|
"//tensorflow/python:math_ops",
|
|
"//tensorflow/python:training",
|
|
"//tensorflow/python:variables",
|
|
"@absl_py//absl/testing:parameterized",
|
|
],
|
|
)
|
|
|
|
py_binary(
|
|
name = "inspect_checkpoint",
|
|
srcs = ["inspect_checkpoint.py"],
|
|
python_version = "PY3",
|
|
srcs_version = "PY2AND3",
|
|
deps = [":inspect_checkpoint_lib"],
|
|
)
|
|
|
|
py_library(
|
|
name = "inspect_checkpoint_lib",
|
|
srcs = ["inspect_checkpoint.py"],
|
|
srcs_version = "PY2AND3",
|
|
deps = [
|
|
"//tensorflow/python", # TODO(b/34059704): remove when fixed
|
|
"//tensorflow/python:platform",
|
|
"//tensorflow/python:py_checkpoint_reader",
|
|
],
|
|
)
|
|
|
|
py_library(
|
|
name = "strip_unused_lib",
|
|
srcs = ["strip_unused_lib.py"],
|
|
srcs_version = "PY2AND3",
|
|
deps = [
|
|
"//tensorflow/core:protos_all_py",
|
|
"//tensorflow/python:framework",
|
|
"//tensorflow/python:platform",
|
|
],
|
|
)
|
|
|
|
py_library(
|
|
name = "module_util",
|
|
srcs = ["module_util.py"],
|
|
srcs_version = "PY2AND3",
|
|
)
|
|
|
|
py_binary(
|
|
name = "strip_unused",
|
|
srcs = ["strip_unused.py"],
|
|
python_version = "PY3",
|
|
srcs_version = "PY2AND3",
|
|
deps = [
|
|
":strip_unused_lib",
|
|
"//tensorflow/python:framework_for_generated_wrappers",
|
|
"//tensorflow/python:platform",
|
|
"@six_archive//:six",
|
|
],
|
|
)
|
|
|
|
py_test(
|
|
name = "strip_unused_test",
|
|
size = "small",
|
|
srcs = ["strip_unused_test.py"],
|
|
python_version = "PY3",
|
|
srcs_version = "PY2AND3",
|
|
tags = ["notap"],
|
|
deps = [
|
|
":strip_unused_lib",
|
|
"//tensorflow/core:protos_all_py",
|
|
"//tensorflow/python:client",
|
|
"//tensorflow/python:client_testlib",
|
|
"//tensorflow/python:framework",
|
|
"//tensorflow/python:framework_for_generated_wrappers",
|
|
"//tensorflow/python:framework_test_lib",
|
|
"//tensorflow/python:math_ops",
|
|
],
|
|
)
|
|
|
|
py_library(
|
|
name = "optimize_for_inference_lib",
|
|
srcs = ["optimize_for_inference_lib.py"],
|
|
srcs_version = "PY2AND3",
|
|
deps = [
|
|
":strip_unused_lib",
|
|
"//tensorflow/core:protos_all_py",
|
|
"//tensorflow/python:framework",
|
|
"//tensorflow/python:framework_for_generated_wrappers",
|
|
"//tensorflow/python:platform",
|
|
"//third_party/py/numpy",
|
|
"@six_archive//:six",
|
|
],
|
|
)
|
|
|
|
py_binary(
|
|
name = "optimize_for_inference",
|
|
srcs = ["optimize_for_inference.py"],
|
|
python_version = "PY3",
|
|
srcs_version = "PY2AND3",
|
|
deps = [":optimize_for_inference_main_lib"],
|
|
)
|
|
|
|
py_library(
|
|
name = "optimize_for_inference_main_lib",
|
|
srcs = ["optimize_for_inference.py"],
|
|
srcs_version = "PY2AND3",
|
|
deps = [
|
|
":optimize_for_inference_lib",
|
|
"//tensorflow/core:protos_all_py",
|
|
"//tensorflow/python", # TODO(b/34059704): remove when fixed
|
|
"//tensorflow/python:framework",
|
|
"//tensorflow/python:framework_for_generated_wrappers",
|
|
"//tensorflow/python:platform",
|
|
"@six_archive//:six",
|
|
],
|
|
)
|
|
|
|
py_test(
|
|
name = "optimize_for_inference_test",
|
|
size = "small",
|
|
srcs = ["optimize_for_inference_test.py"],
|
|
python_version = "PY3",
|
|
srcs_version = "PY2AND3",
|
|
deps = [
|
|
":optimize_for_inference_lib",
|
|
"//tensorflow/core:protos_all_py",
|
|
"//tensorflow/python:array_ops",
|
|
"//tensorflow/python:client_testlib",
|
|
"//tensorflow/python:framework",
|
|
"//tensorflow/python:framework_for_generated_wrappers",
|
|
"//tensorflow/python:image_ops",
|
|
"//tensorflow/python:math_ops",
|
|
"//tensorflow/python:nn_ops",
|
|
"//tensorflow/python:nn_ops_gen",
|
|
"//third_party/py/numpy",
|
|
],
|
|
)
|
|
|
|
py_library(
|
|
name = "selective_registration_header_lib",
|
|
srcs = ["selective_registration_header_lib.py"],
|
|
srcs_version = "PY2AND3",
|
|
visibility = ["//visibility:public"],
|
|
deps = [
|
|
"//tensorflow/python", # TODO(b/34059704): remove when fixed
|
|
"//tensorflow/python:_pywrap_kernel_registry",
|
|
"//tensorflow/python:platform",
|
|
],
|
|
)
|
|
|
|
py_binary(
|
|
name = "print_selective_registration_header",
|
|
srcs = ["print_selective_registration_header.py"],
|
|
python_version = "PY3",
|
|
srcs_version = "PY2AND3",
|
|
visibility = ["//visibility:public"],
|
|
deps = [":print_selective_registration_header_lib"],
|
|
)
|
|
|
|
py_library(
|
|
name = "print_selective_registration_header_lib",
|
|
srcs = ["print_selective_registration_header.py"],
|
|
srcs_version = "PY2AND3",
|
|
visibility = ["//visibility:public"],
|
|
deps = [
|
|
":selective_registration_header_lib",
|
|
"//tensorflow/python:platform",
|
|
],
|
|
)
|
|
|
|
py_test(
|
|
name = "print_selective_registration_header_test",
|
|
srcs = ["print_selective_registration_header_test.py"],
|
|
python_version = "PY3",
|
|
srcs_version = "PY2AND3",
|
|
deps = [
|
|
":selective_registration_header_lib",
|
|
"//tensorflow/python:client_testlib",
|
|
"//tensorflow/python:platform",
|
|
],
|
|
)
|
|
|
|
py_binary(
|
|
name = "saved_model_cli",
|
|
srcs = ["saved_model_cli.py"],
|
|
python_version = "PY3",
|
|
srcs_version = "PY2AND3",
|
|
deps = [":saved_model_cli_lib"],
|
|
)
|
|
|
|
py_library(
|
|
name = "saved_model_cli_lib",
|
|
srcs = ["saved_model_cli.py"],
|
|
srcs_version = "PY2AND3",
|
|
deps = [
|
|
# Note: if you make any changes here, make corresponding changes to the
|
|
# deps of the "tools_pip" target in this file. Otherwise release builds
|
|
# (built with --define=no_tensorflow_py_deps=true) may end up with a
|
|
# broken saved_model_cli.
|
|
":saved_model_aot_compile",
|
|
":saved_model_utils",
|
|
"//tensorflow/python",
|
|
"//tensorflow/python/debug:local_cli_wrapper",
|
|
],
|
|
)
|
|
|
|
py_library(
|
|
name = "saved_model_aot_compile",
|
|
srcs = ["saved_model_aot_compile.py"],
|
|
srcs_version = "PY2AND3",
|
|
deps = [
|
|
"//tensorflow/python",
|
|
"//tensorflow/python:tf_optimizer",
|
|
] + if_xla_available(
|
|
["//tensorflow/compiler/tf2xla:tf2xla_proto_py"],
|
|
),
|
|
)
|
|
|
|
py_test(
|
|
name = "saved_model_cli_test",
|
|
srcs = ["saved_model_cli_test.py"],
|
|
data = [
|
|
"//tensorflow/cc/saved_model:saved_model_half_plus_two",
|
|
],
|
|
python_version = "PY3",
|
|
srcs_version = "PY2AND3",
|
|
tags = [
|
|
"manual",
|
|
"no-internal-py3",
|
|
"nosan",
|
|
],
|
|
deps = [
|
|
":saved_model_cli_lib",
|
|
"//tensorflow/core:protos_all_py",
|
|
"//tensorflow/python:client_testlib",
|
|
"@absl_py//absl/testing:parameterized",
|
|
],
|
|
)
|
|
|
|
saved_model_compile_aot(
|
|
name = "aot_compiled_x_plus_y",
|
|
cpp_class = "XPlusY",
|
|
directory = "//tensorflow/cc/saved_model:testdata/x_plus_y_v2_debuginfo",
|
|
filegroups = [
|
|
"//tensorflow/cc/saved_model:saved_model_half_plus_two",
|
|
],
|
|
force_without_xla_support_flag = False,
|
|
tags = ["no_rocm"],
|
|
)
|
|
|
|
saved_model_compile_aot(
|
|
name = "aot_compiled_vars_and_arithmetic_frozen",
|
|
cpp_class = "VarsAndArithmeticFrozen",
|
|
directory = "//tensorflow/cc/saved_model:testdata/VarsAndArithmeticObjectGraph",
|
|
filegroups = [
|
|
"//tensorflow/cc/saved_model:saved_model_half_plus_two",
|
|
],
|
|
force_without_xla_support_flag = False,
|
|
tags = ["no_rocm"],
|
|
)
|
|
|
|
saved_model_compile_aot(
|
|
name = "aot_compiled_vars_and_arithmetic",
|
|
cpp_class = "VarsAndArithmetic",
|
|
directory = "//tensorflow/cc/saved_model:testdata/VarsAndArithmeticObjectGraph",
|
|
filegroups = [
|
|
"//tensorflow/cc/saved_model:saved_model_half_plus_two",
|
|
],
|
|
force_without_xla_support_flag = False,
|
|
tags = ["no_rocm"],
|
|
variables_to_feed = "variable_x",
|
|
)
|
|
|
|
tf_cc_test(
|
|
name = "aot_compiled_test",
|
|
srcs = if_xla_available([
|
|
"aot_compiled_test.cc",
|
|
]),
|
|
tags = ["no_rocm"],
|
|
deps = [
|
|
"//tensorflow/core:test_main",
|
|
] + if_xla_available([
|
|
":aot_compiled_vars_and_arithmetic",
|
|
":aot_compiled_vars_and_arithmetic_frozen",
|
|
":aot_compiled_x_plus_y",
|
|
"//tensorflow/core:test",
|
|
"//tensorflow/core/platform:logging",
|
|
]),
|
|
)
|