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/python/grappler/BUILD
+707
View File
@@ -0,0 +1,707 @@
load("@rules_cc//cc:cc_library.bzl", "cc_library")
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:tensorflow.bzl", "if_not_windows")
load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test", "get_compatible_with_portable", "tf_py_strict_test", "tf_pybind_cc_library_wrapper", "tf_python_pybind_extension")
load("//tensorflow/core/platform:build_config.bzl", "tf_protos_grappler")
load("//tensorflow/core/platform:build_config_root.bzl", "if_pywrap")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = ["//tensorflow:internal"],
licenses = ["notice"],
)
# TODO(gunan): Investigate making this action hermetic so we do not need
# to run it locally.
cc_library(
name = "cost_analyzer_lib",
srcs = ["cost_analyzer.cc"],
hdrs = ["cost_analyzer.h"],
compatible_with = get_compatible_with_portable(),
deps = [
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/grappler:grappler_item",
"//tensorflow/core/grappler/clusters:cluster",
"//tensorflow/core/grappler/costs:analytical_cost_estimator",
"//tensorflow/core/grappler/costs:cost_estimator",
"//tensorflow/core/grappler/costs:measuring_cost_estimator",
"//tensorflow/core/grappler/costs:utils",
"@com_google_absl//absl/log",
"@com_google_absl//absl/status",
] + tf_protos_grappler(),
alwayslink = 1,
)
# Necessary for the pywrap inclusion below. Combining targets does not work
# properly.
tf_pybind_cc_library_wrapper(
name = "cost_analyzer_headers",
deps = [
":cost_analyzer_lib",
],
)
tf_python_pybind_extension(
name = "_pywrap_cost_analyzer",
srcs = ["cost_analyzer_wrapper.cc"],
hdrs = [
"cost_analyzer.h",
"//tensorflow/cc:pywrap_required_hdrs",
"//tensorflow/core/grappler:pywrap_required_hdrs",
"//tensorflow/core/grappler/clusters:pywrap_required_hdrs",
"//tensorflow/core/grappler/costs:pywrap_required_hdrs",
"//tensorflow/core/public:session.h",
"//tensorflow/core/public:session_options.h",
],
enable_stub_generation = True,
pytype_srcs = [
"_pywrap_cost_analyzer.pyi",
],
starlark_only = True,
deps = [
":cost_analyzer_headers",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
"//tensorflow/core:framework_headers_lib",
"//tensorflow/core:framework_lite",
"//tensorflow/core:lib",
"//tensorflow/core:lib_headers_for_pybind",
"//tensorflow/core:lib_proto_parsing",
"//tensorflow/core:portable_gif_internal",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/common_runtime:core_cpu_headers_lib",
"//tensorflow/core/common_runtime:device_set",
"//tensorflow/core/common_runtime/gpu:gpu_id",
"//tensorflow/core/framework:allocator",
"//tensorflow/core/framework:tensor",
"//tensorflow/core/platform:threadpool_options",
"//tensorflow/python/lib/core:pybind11_status",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@pybind11",
"@tsl//tsl/platform:platform_port",
"@tsl//tsl/platform:thread_annotations",
"@xla//xla/tsl/protobuf:error_codes_proto_impl_cc",
],
)
cc_library(
name = "model_analyzer_lib",
srcs = ["model_analyzer.cc"],
hdrs = ["model_analyzer.h"],
deps = [
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/grappler:grappler_item",
"//tensorflow/core/grappler/costs:graph_properties",
"@com_google_absl//absl/status",
],
)
tf_python_pybind_extension(
name = "_pywrap_model_analyzer",
srcs = ["model_analyzer_wrapper.cc"],
hdrs = [
"model_analyzer.h",
"//tensorflow/core/grappler:pywrap_required_hdrs",
],
enable_stub_generation = True,
pytype_srcs = [
"_pywrap_model_analyzer.pyi",
],
starlark_only = True,
deps = [
"//tensorflow/core:framework_headers_lib",
"//tensorflow/core:lib",
"//tensorflow/core:lib_headers_for_pybind",
"//tensorflow/core:lib_proto_parsing",
"//tensorflow/core:portable_gif_internal",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/framework:tensor",
"//tensorflow/python/lib/core:pybind11_status",
"@com_google_absl//absl/status",
"@pybind11",
"@tsl//tsl/platform:platform_port",
] + if_pywrap(["//tensorflow/python/grappler:model_analyzer_lib"]),
)
py_library(
name = "tf_item",
srcs = ["item.py"],
strict_deps = True,
visibility = ["//visibility:public"],
deps = [
":_pywrap_tf_item",
"//tensorflow/core:protos_all_py",
"//tensorflow/core/grappler/costs:op_performance_data_py",
],
)
tf_python_pybind_extension(
name = "_pywrap_tf_item",
srcs = ["item_wrapper.cc"],
hdrs = [
"//tensorflow/cc:pywrap_required_hdrs",
"//tensorflow/core/grappler:pywrap_required_hdrs",
"//tensorflow/core/grappler/clusters:pywrap_required_hdrs",
"//tensorflow/core/grappler/costs:pywrap_required_hdrs",
"//tensorflow/core/grappler/utils:pywrap_required_hdrs",
],
enable_stub_generation = True,
pytype_srcs = [
"_pywrap_tf_item.pyi",
],
deps = [
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
"//tensorflow/core:framework_headers_lib",
"//tensorflow/core:framework_lite",
"//tensorflow/core:lib",
"//tensorflow/core:lib_proto_parsing",
"//tensorflow/core:portable_gif_internal",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:session_options",
"//tensorflow/core/common_runtime:core_cpu_headers_lib",
"//tensorflow/core/common_runtime:device_set",
"//tensorflow/core/common_runtime/gpu:gpu_id",
"//tensorflow/core/framework:allocator",
"//tensorflow/core/framework:tensor",
"//tensorflow/core/grappler:utils",
"//tensorflow/python/lib/core:pybind11_status",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/types:span",
"@pybind11",
"@tsl//tsl/platform:platform_port",
"@tsl//tsl/platform:thread_annotations",
"@xla//xla/tsl/protobuf:error_codes_proto_impl_cc",
] + if_not_windows(["//tensorflow/core/grappler/costs:graph_properties"]), # b/148556093,
)
tf_py_strict_test(
name = "item_test",
size = "small",
srcs = ["item_test.py"],
tags = [
"grappler",
"no_pip", # tf_optimizer is not available in pip.
],
deps = [
":tf_item",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:meta_graph",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops_gen",
"//tensorflow/python/ops:control_flow_ops",
"//tensorflow/python/ops:state_ops",
"//tensorflow/python/ops:variable_v1",
"//tensorflow/python/platform:client_testlib",
],
)
tf_py_strict_test(
name = "datasets_test",
size = "medium",
srcs = ["datasets_test.py"],
tags = [
"grappler",
"no_pip", # tf_optimizer is not available in pip.
],
deps = [
":tf_item",
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/data/ops:iterator_ops",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:meta_graph",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
],
)
py_library(
name = "tf_cluster",
srcs = ["cluster.py"],
strict_deps = True,
visibility = ["//visibility:public"],
deps = [
":_pywrap_tf_cluster",
"//tensorflow/core:protos_all_py",
"//tensorflow/core/grappler/costs:op_performance_data_py",
],
)
tf_python_pybind_extension(
name = "_pywrap_tf_cluster",
srcs = ["cluster_wrapper.cc"],
hdrs = [
"//tensorflow/cc:pywrap_required_hdrs",
] + if_pywrap(
if_false = [
"//tensorflow/core/grappler:pywrap_required_hdrs",
"//tensorflow/core/grappler/clusters:pywrap_required_hdrs",
"//tensorflow/core/grappler/costs:pywrap_required_hdrs",
"//tensorflow/core/grappler/utils:pywrap_required_hdrs",
],
),
enable_stub_generation = True,
pytype_srcs = [
"_pywrap_tf_cluster.pyi",
],
deps = [
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
"//tensorflow/core:framework_headers_lib",
"//tensorflow/core:framework_lite",
"//tensorflow/core:lib",
"//tensorflow/core:lib_headers_for_pybind",
"//tensorflow/core:lib_proto_parsing",
"//tensorflow/core:portable_gif_internal",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:session_options",
"//tensorflow/core/common_runtime:core_cpu_headers_lib",
"//tensorflow/core/common_runtime:device_set",
"//tensorflow/core/common_runtime/gpu:gpu_id",
"//tensorflow/core/framework:allocator",
"//tensorflow/core/framework:tensor",
"//tensorflow/core/grappler:utils",
"//tensorflow/core/platform:status",
"//tensorflow/python/lib/core:pybind11_status",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/types:span",
"@pybind11",
"@tsl//tsl/platform:platform_port",
"@tsl//tsl/platform:thread_annotations",
"@xla//xla/tsl/protobuf:error_codes_proto_impl_cc",
] + if_pywrap(
if_true = [
"//tensorflow/core/grappler/costs:measuring_cost_estimator",
"//tensorflow/core/grappler/clusters:single_machine",
],
),
)
cuda_py_strict_test(
name = "cluster_test",
size = "small",
srcs = ["cluster_test.py"],
shard_count = 5,
tags = [
"grappler",
"no_pip", # tf_optimizer is not available in pip.
"no_windows", # b/173520599
"notap", # TODO(b/135924227): Re-enable after fixing flakiness.
],
# This test will not run on XLA because it primarily tests the TF Classic flow.
xla_enable_strict_auto_jit = False,
deps = [
":tf_cluster",
":tf_item",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/framework:meta_graph",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:random_ops",
"//tensorflow/python/platform:client_testlib",
],
)
py_library(
name = "tf_optimizer",
srcs = ["tf_optimizer.py"],
strict_deps = True,
visibility = ["//visibility:public"],
deps = [
":_pywrap_tf_optimizer",
":tf_cluster",
"//tensorflow/core:protos_all_py",
"@absl_py//absl/logging",
],
)
tf_python_pybind_extension(
name = "_pywrap_tf_optimizer",
srcs = ["tf_optimizer_wrapper.cc"],
hdrs = if_pywrap(
if_false = [
"//tensorflow/cc:pywrap_required_hdrs",
"//tensorflow/core/grappler:pywrap_required_hdrs",
"//tensorflow/core/grappler/clusters:pywrap_required_hdrs",
"//tensorflow/core/grappler/costs:pywrap_required_hdrs",
"//tensorflow/core/grappler/optimizers:pywrap_required_hdrs",
"//tensorflow/core/grappler/verifiers:pywrap_required_hdrs",
],
),
enable_stub_generation = True,
pytype_srcs = [
"_pywrap_tf_optimizer.pyi",
],
# This fails Windows builds. Please check b/266870200 for details.
# dynamic_deps = ["//tensorflow/python:_pywrap_tensorflow_internal.so"] + select({
# "//tensorflow:macos": ["//tensorflow:libtensorflow_framework.%s.dylib" % VERSION],
# "//conditions:default": ["//tensorflow:libtensorflow_framework.so.%s" % VERSION],
# "//tensorflow:windows": [],
# }),
# static_deps = tf_python_pybind_static_deps(),
deps = [
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
"//tensorflow/core:framework_headers_lib",
"//tensorflow/core:framework_lite",
"//tensorflow/core:lib",
"//tensorflow/core:lib_headers_for_pybind",
"//tensorflow/core:lib_proto_parsing",
"//tensorflow/core:portable_gif_internal",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:session_options",
"//tensorflow/core/common_runtime:core_cpu_headers_lib",
"//tensorflow/core/common_runtime:device",
"//tensorflow/core/common_runtime:device_factory",
"//tensorflow/core/common_runtime:device_set",
"//tensorflow/core/common_runtime/gpu:gpu_id",
"//tensorflow/core/framework:allocator",
"//tensorflow/core/framework:tensor",
"//tensorflow/core/platform:errors",
"//tensorflow/core/platform:status",
"//tensorflow/python/lib/core:pybind11_status",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@pybind11",
"@pybind11_protobuf//pybind11_protobuf:native_proto_caster",
"@tsl//tsl/platform:platform_port",
"@tsl//tsl/platform:thread_annotations",
"@xla//xla/tsl/protobuf:error_codes_proto_impl_cc",
] + if_pywrap(
if_true = [
"//tensorflow/core/grappler/clusters:cluster",
"//tensorflow/core/grappler/clusters:utils",
"//tensorflow/core/grappler:grappler_item_builder",
"//tensorflow/core/grappler/optimizers:meta_optimizer",
"//tensorflow/core/grappler/optimizers:graph_optimizer",
"//tensorflow/core/grappler/verifiers:graph_verifier",
],
),
)
tf_py_strict_test(
name = "tf_optimizer_test",
size = "small",
srcs = ["tf_optimizer_test.py"],
tags = [
"grappler",
"no_pip", # tf_optimizer is not available in pip.
],
deps = [
":tf_item",
":tf_optimizer",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:meta_graph",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:variable_v1",
"//tensorflow/python/ops:while_loop",
"//tensorflow/python/platform:client_testlib",
],
)
tf_py_strict_test(
name = "memory_optimizer_test",
size = "medium",
srcs = ["memory_optimizer_test.py"],
tags = [
"grappler",
],
deps = [
":tf_optimizer",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/client:session",
"//tensorflow/python/framework:meta_graph",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:random_seed",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:nn",
"//tensorflow/python/ops:variable_scope",
"//tensorflow/python/ops:variable_v1",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/training:training_lib",
],
)
cuda_py_strict_test(
name = "constant_folding_test",
size = "medium",
srcs = ["constant_folding_test.py"],
tags = [
"grappler",
],
deps = [
"//tensorflow/python/eager:backprop",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:functional_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/ops:while_loop",
"//tensorflow/python/platform:client_testlib",
"//third_party/py/numpy",
],
)
cuda_py_strict_test(
name = "arithmetic_optimizer_test",
size = "small",
srcs = ["arithmetic_optimizer_test.py"],
tags = [
"grappler",
],
xla_enable_strict_auto_jit = False,
deps = [
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/platform:client_testlib",
],
)
# TODO(b/131764887) Remove once LayoutOptimizer is swapped out with GenericLayoutOptimizer.
#
# cuda_py_test(
# name = "layout_optimizer_test",
# size = "medium",
# srcs = [
# "layout_optimizer_test.py",
# ],
# deps = [
# "//tensorflow/python/platform:client_testlib",
# "//tensorflow/python/framework:for_generated_wrappers",
# "//tensorflow/python/ops:array_ops",
# "//tensorflow/python/ops:functional_ops",
# "//tensorflow/python/ops:math_ops",
# "//tensorflow/python:nn",
# "//tensorflow/python/user_ops:ops",
# "//tensorflow/python/ops:random_ops",
# "//tensorflow/python/ops:state_ops",
# ":tf_cluster",
# ":tf_optimizer",
# "//tensorflow/python/training:training",
# "//third_party/py/numpy",
# "//tensorflow/core:protos_all_py",
# "//tensorflow/python/framework:constant_op",
# "//tensorflow/python/framework:dtypes",
# ],
# shard_count = 10,
# tags = [
# "grappler",
# ],
# # This test will not run on XLA because it primarily tests the TF Classic flow.
# xla_enable_strict_auto_jit = False,
# )
py_library(
name = "cost_analyzer",
srcs = ["cost_analyzer.py"],
strict_deps = True,
deps = [
":_pywrap_cost_analyzer",
":tf_cluster",
":tf_item",
],
)
py_binary(
name = "cost_analyzer_tool",
srcs = ["cost_analyzer_tool.py"],
strict_deps = True,
deps = [
":cost_analyzer",
":tf_optimizer",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/framework:importer",
"//tensorflow/python/framework:ops",
"//tensorflow/python/platform:gfile",
"//tensorflow/python/training:saver",
"@absl_py//absl:app",
],
)
tf_py_strict_test(
name = "cost_analyzer_test",
size = "small",
srcs = ["cost_analyzer_test.py"],
tags = [
"grappler",
"no_cuda_on_cpu_tap",
"no_mac",
"no_pip",
"no_windows", # TODO(b/151942037)
],
deps = [
":cost_analyzer",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:meta_graph",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:nn_grad",
"//tensorflow/python/ops:nn_ops",
"//tensorflow/python/ops:random_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/training:adam",
],
)
py_library(
name = "model_analyzer",
srcs = [
"model_analyzer.py",
],
strict_deps = True,
deps = [":_pywrap_model_analyzer"],
)
tf_py_strict_test(
name = "model_analyzer_test",
size = "small",
srcs = ["model_analyzer_test.py"],
tags = [
"grappler",
"no_pip",
],
deps = [
":model_analyzer",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:meta_graph",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/platform:client_testlib",
],
)
cuda_py_strict_test(
name = "auto_mixed_precision_test",
size = "medium",
srcs = [
"auto_mixed_precision_test.py",
],
tags = ["grappler"],
# This test analyzes the graph, but XLA changes the names of nodes.
xla_enable_strict_auto_jit = False,
deps = [
"//tensorflow/core:protos_all_py",
"//tensorflow/python:tf2",
"//tensorflow/python/client:session",
"//tensorflow/python/data/ops:dataset_ops",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:function",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:random_seed",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:init_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:nn",
"//tensorflow/python/ops:nn_impl",
"//tensorflow/python/ops:random_ops",
"//tensorflow/python/ops:tensor_array_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/ops:while_loop",
"//tensorflow/python/ops/losses",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/platform:sysconfig",
"//tensorflow/python/training:adam",
"//tensorflow/python/training:gradient_descent",
"//tensorflow/python/util:_pywrap_utils",
"//third_party/py/numpy",
"@absl_py//absl/testing:parameterized",
],
)
cuda_py_strict_test(
name = "remapper_test",
size = "medium",
srcs = ["remapper_test.py"],
tags = ["grappler"],
# This test analyzes the graph, but XLA changes the names of nodes.
xla_enable_strict_auto_jit = False,
deps = [
":tf_optimizer",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/client:session",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:meta_graph",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:init_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:nn",
"//tensorflow/python/ops:nn_ops",
"//tensorflow/python/ops:random_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/platform:sysconfig",
"//tensorflow/python/util:_pywrap_utils",
"@absl_py//absl/testing:parameterized",
],
)
tf_python_pybind_extension(
name = "_pywrap_graph_analyzer",
srcs = ["graph_analyzer_tool_wrapper.cc"],
enable_stub_generation = True,
pytype_srcs = [
"_pywrap_graph_analyzer.pyi",
],
deps = [
"//tensorflow/core/grappler/graph_analyzer:graph_analyzer_tool",
"@pybind11",
],
)
py_binary(
name = "graph_analyzer",
srcs = ["graph_analyzer.py"],
strict_deps = True,
deps = [
":_pywrap_graph_analyzer",
"@absl_py//absl:app",
],
)
tensorflow/python/grappler/_pywrap_cost_analyzer.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 GenerateCostReport(arg0: bytes, arg1: bool, arg2: bool, arg3) -> bytes: ...
tensorflow/python/grappler/_pywrap_graph_analyzer.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 GraphAnalyzer(arg0: str, arg1: int) -> None: ...
tensorflow/python/grappler/_pywrap_model_analyzer.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 GenerateModelReport(arg0: bytes, arg1: bool, arg2: bool) -> bytes: ...
tensorflow/python/grappler/_pywrap_tf_cluster.pyi
+27
View File
@@ -0,0 +1,27 @@
# 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 Cluster:
def __init__(self, *args, **kwargs) -> None: ...
def TF_DeterminePeakMemoryUsage(arg0, arg1: Cluster) -> dict[str, tuple[int, list[tuple[str, int, int, int, int]]]]: ...
def TF_EstimatePerformance(arg0: bytes) -> float: ...
def TF_GetSupportedDevices(arg0: Cluster, arg1) -> dict[str, list[str]]: ...
def TF_ListAvailableOps() -> list[str]: ...
def TF_ListDevices(arg0: Cluster) -> list[bytes]: ...
def TF_MeasureCosts(arg0, arg1: Cluster, arg2: bool) -> tuple[list[bytes], float, bytes]: ...
def TF_NewCluster(arg0: bool, arg1: bool) -> Cluster: ...
def TF_NewVirtualCluster(arg0: list[bytes]) -> Cluster: ...
def TF_ShutdownCluster(arg0: Cluster) -> None: ...
tensorflow/python/grappler/_pywrap_tf_item.pyi
+22
View File
@@ -0,0 +1,22 @@
# 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 GrapplerItem:
def __init__(self, *args, **kwargs) -> None: ...
def TF_GetColocationGroups(arg0: GrapplerItem) -> list[list[str]]: ...
def TF_GetOpProperties(arg0: GrapplerItem) -> dict[str, list[bytes]]: ...
def TF_IdentifyImportantOps(arg0: GrapplerItem, arg1: bool) -> list[str]: ...
def TF_NewItem(arg0: bytes, arg1: bool, arg2: bool) -> GrapplerItem: ...
tensorflow/python/grappler/_pywrap_tf_optimizer.pyi
+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_OptimizeGraph(*args, **kwargs): ...
def TF_OptimizeGraphSerialized(arg0, arg1: str, arg2: str, arg3: bool, arg4: str, arg5: bool) -> bytes: ...
tensorflow/python/grappler/arithmetic_optimizer_test.py
+46
View File
@@ -0,0 +1,46 @@
# 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 Grappler Arithmetic Optimizer."""
from tensorflow.python.eager import context
from tensorflow.python.eager import def_function
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.platform import test
class ArithmeticOptimizerTest(test.TestCase):
# See b/146524878.
def testFunctionArgShapeInference(self):
@def_function.function
def f(x, y):
return math_ops.matmul(
x, array_ops.reshape(array_ops.transpose(y), [384, 1536]))
with context.eager_mode():
x = array_ops.ones((1, 384))
y = array_ops.ones((1536, 384))
with context.collect_graphs(optimized=True) as graphs:
f(x, y).numpy()
self.assertLen(graphs, 1)
self.assertLen(graphs[0].node, 4)
self.assertEqual(graphs[0].node[2].name,
'ArithmeticOptimizer/FoldTransposeIntoMatMul_MatMul')
if __name__ == '__main__':
test.main()
tensorflow/python/grappler/auto_mixed_precision_test.py
+906
View File
@@ -0,0 +1,906 @@
# 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 Grappler AutoMixedPrecision."""
import os
import re
from absl.testing import parameterized
import numpy as np
from tensorflow.core.framework import types_pb2
from tensorflow.core.protobuf import config_pb2
from tensorflow.core.protobuf import rewriter_config_pb2
from tensorflow.python import tf2
from tensorflow.python.client import session
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import function
from tensorflow.python.framework import ops
from tensorflow.python.framework import random_seed
from tensorflow.python.framework import test_util
from tensorflow.python.layers import layers
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import init_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nn
from tensorflow.python.ops import nn_impl
from tensorflow.python.ops import random_ops
from tensorflow.python.ops import tensor_array_ops
from tensorflow.python.ops import variables
from tensorflow.python.ops import while_loop
from tensorflow.python.ops.losses import losses
from tensorflow.python.platform import sysconfig as sysconfig_lib
from tensorflow.python.platform import test
from tensorflow.python.training import adam
from tensorflow.python.training import gradient_descent
from tensorflow.python.util import _pywrap_utils
def _input(shape):
"""Generates an input of a given shape."""
return variables.Variable(random_ops.truncated_normal(shape, seed=0))
def _weight(shape):
"""Generates a weight of a given shape."""
# Note that the lambda is needed to allow construction inside loops.
return variables.Variable(lambda: init_ops.glorot_uniform_initializer(seed=0)
(shape))
def _bias(shape):
"""Generates a bias of a given shape."""
return constant_op.constant(0.1, shape=shape)
def _conv2d(x, w):
"""Returns a 2d convolution layer with full stride."""
return nn.conv2d(x, w, strides=[1, 1, 1, 1], padding='SAME')
def _conv3d(x, w):
"""Returns a 3d convolution layer with full stride."""
return nn.conv3d(x, w, strides=[1, 1, 1, 1, 1], padding='SAME')
def _max_pool_2x2(x):
"""Downsamples a feature map by 2X."""
return nn.max_pool(
x, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
def _fused_batchnorm(x, scale, offset):
"""Batchnorm."""
return nn_impl.fused_batch_norm(
x, scale=scale, offset=offset, is_training=True)
def _conv_bn(x):
"""Conv followed by batchnorm."""
i = array_ops.reshape(x, [-1, 8, 8, 1])
f = _weight([3, 3, 1, 6])
x = _conv2d(i, f)
s = _weight([6])
o = _weight([6])
y, _, _ = _fused_batchnorm(x, s, o)
y = array_ops.identity(y)
return y
def _conv3d_bn(x):
"""Conv3D followed by batchnorm."""
i = array_ops.reshape(x, [-1, 8, 8, 8, 1])
f = _weight([3, 3, 3, 1, 6])
x = _conv3d(i, f)
s = _weight([6])
o = _weight([6])
x = array_ops.reshape(x, [-1, 8, 8, 6])
y, _, _ = _fused_batchnorm(x, s, o)
y = array_ops.identity(y)
return y
def _matmul_act(x):
"""Matmul followed by activation."""
i = array_ops.reshape(x, [8, 8])
f = _weight([8, 8])
x = math_ops.matmul(i, f)
y = nn.relu(x)
return y
def _conv_pool(x):
"""(Conv -> bias -> relu -> max_pool) x2."""
x_image = array_ops.reshape(x, [-1, 8, 8, 1])
w_conv1 = _weight([3, 3, 1, 6])
b_conv1 = _bias([6])
h_conv1 = nn.relu(nn.bias_add(_conv2d(x_image, w_conv1), b_conv1))
h_pool1 = _max_pool_2x2(h_conv1)
w_conv2 = _weight([3, 3, 6, 4])
b_conv2 = _bias([4])
h_conv2 = nn.relu(nn.bias_add(_conv2d(h_pool1, w_conv2), b_conv2))
h_pool2 = _max_pool_2x2(h_conv2)
return h_pool2
def _depthwise_conv2d(x, w):
"""Returns a 2d depthwise convolution layer with full stride."""
return nn.depthwise_conv2d(x, w, strides=[1, 1, 1, 1], padding='SAME')
def _simple_loop(x, functor):
"""Simple loop whose body is provided by the functor."""
init = (constant_op.constant(0), x)
c = lambda i, j: i < 4
b = lambda i, j: (i + 1, functor(j))
ij = while_loop.while_loop(c, b, init)
return ij
def _loop_vars_intertwined(x0, y0, functor_x, functor_y):
"""Loop whose loop variables are intertwined."""
c = lambda i, j, x, y: j < 4
b = lambda i, j, x, y: (j + 1, i + 1, functor_y(y), functor_x(x))
init = (constant_op.constant(0), constant_op.constant(0), x0, y0)
ijzw = while_loop.while_loop(c, b, init)
return ijzw
def _lstm_cell(prev_c, prev_h, x):
"""Create an LSTM cell."""
# i: input gate
# f: forget gate
# o: output gate
# c: cell state
# x: input
# h: embedding
bias = _bias([4])
w = _weight([8, 16])
ifoc = math_ops.matmul(array_ops.concat([x, prev_h], axis=1), w)
i, f, o, c = array_ops.split(ifoc, 4, axis=1)
i = math_ops.sigmoid(nn.bias_add(i, bias))
f = math_ops.sigmoid(nn.bias_add(f, bias))
o = math_ops.sigmoid(nn.bias_add(o, bias))
c = math_ops.tanh(nn.bias_add(c, bias))
next_c = f * prev_c + i * c
next_h = o * math_ops.tanh(next_c)
return next_c, next_h
def _recurrent_lstm(c, h):
"""Dynamic single-layer LSTM with TensorArray."""
def cond(i, c, h, ta_x):
del c, h, ta_x
return i < 4
def body(i, c, h, ta_x):
x = ta_x.read(i)
next_c, next_h = _lstm_cell(c, h, x)
return (i + 1, next_c, next_h, ta_x)
ta_x = tensor_array_ops.TensorArray(dtype=dtypes.float32, size=4)
for i in range(0, 4):
ta_x = ta_x.write(
i, constant_op.constant(0.1, shape=[8, 4], dtype=dtypes.float32))
init = (constant_op.constant(0), c, h, ta_x)
r = while_loop.while_loop(cond, body, init)
return r
def _make_node_with_color(color, input_tensor, name=None):
"""Returns a node representative of the specified list type."""
color = color.lower()
if color == 'w': # Allow node
weights = _weight(input_tensor.get_shape().as_list())
return math_ops.matmul(input_tensor, weights, name=name)
if color == 'g': # Infer node
return math_ops.add(input_tensor, 0.1, name=name)
if color == 'c': # Clear node
return nn.relu(input_tensor, name=name)
if color == 'b': # Deny node
return math_ops.pow(math_ops.pow(input_tensor, 2.), 0.5, name=name)
raise ValueError('Invalid node color: ' + str(color))
def _build_simple_loop_graph(inp_colors, body_colors, out_colors):
"""Builds a test graph with a simple loop."""
a = _input([8, 8])
for i, color in enumerate(inp_colors):
a = _make_node_with_color(color, a, 'input_%i' % i)
def body(x):
for i, color in enumerate(body_colors):
x = _make_node_with_color(color, x, 'body_%i' % i)
return x
_, a = _simple_loop(a, body)
for i, color in enumerate(out_colors):
a = _make_node_with_color(color, a, 'output_%i' % i)
a = array_ops.identity(a)
return a
def _get_config(auto_mixed_precision_mode):
"""Returns a ConfigProto with auto mixed precision enabled if appropriate."""
rewrite_config = rewriter_config_pb2.RewriterConfig(
# do not remove duplicated nodes
arithmetic_optimization=rewriter_config_pb2.RewriterConfig.OFF,
# do not turn Conv2D and other nodes into _FusedConv2D
remapping=rewriter_config_pb2.RewriterConfig.OFF,
)
if auto_mixed_precision_mode == 'cuda':
rewrite_config.auto_mixed_precision = rewriter_config_pb2.RewriterConfig.ON
elif auto_mixed_precision_mode == 'mkl':
rewrite_config.auto_mixed_precision_onednn_bfloat16 = (
rewriter_config_pb2.RewriterConfig.ON)
else:
assert auto_mixed_precision_mode is None
rewrite_config.min_graph_nodes = -1
graph_options = config_pb2.GraphOptions(
rewrite_options=rewrite_config, build_cost_model=1)
config = config_pb2.ConfigProto(graph_options=graph_options)
config.graph_options.optimizer_options.opt_level = -1
return config
def _get_device(auto_mixed_precision_mode):
"""Returns the device to run on. If mode is mkl, run on CPU"""
if auto_mixed_precision_mode == 'mkl':
return '/cpu:0'
else:
return ''
def _is_cast_to_fp16(node_name):
return re.match('.*-CastToFp16-[0-9]-AutoMixedPrecision$', node_name)
def _is_cast_to_bf16(node_name):
return re.match('.*-CastToBf16-[0-9]-AutoMixedPrecision$', node_name)
def _is_cast_to_fp32(node_name):
return re.match('.*-CastToFp32-[0-9]-AutoMixedPrecision$', node_name)
def _count_casts(mode, nodes):
"""Counts the number of casts to f16 and fp32."""
num_to_fp16 = 0
num_to_bf16 = 0
num_to_fp32 = 0
for node in nodes:
if _is_cast_to_fp16(node.name):
num_to_fp16 += 1
if _is_cast_to_bf16(node.name):
num_to_bf16 += 1
elif _is_cast_to_fp32(node.name):
num_to_fp32 += 1
if mode == 'cuda':
assert num_to_bf16 == 0
return num_to_fp16, num_to_fp32
else:
assert mode == 'mkl'
assert num_to_fp16 == 0
return num_to_bf16, num_to_fp32
def _build_node_map(nodes):
node_map = {}
for node in nodes:
node_map[node.name] = node
return node_map
def _example_noninlined_funcdef_shape(op):
return [op.inputs[0].shape]
@function.Defun(
shape_func=_example_noninlined_funcdef_shape,
func_name='example_noninlined_funcdef_grad',
noinline=True)
def _example_noninlined_funcdef_grad(features, grad):
"""Gradient of Swish function defined below."""
sigmoid_features = math_ops.sigmoid(features)
activation_grad = (
sigmoid_features * (1.0 + features * (1.0 - sigmoid_features)))
return grad * activation_grad
@function.Defun(
grad_func=_example_noninlined_funcdef_grad,
shape_func=_example_noninlined_funcdef_shape,
func_name='example_noninlined_funcdef',
noinline=True)
def _example_noninlined_funcdef(features):
"""Computes the Swish activation function: `x * sigmoid(x)`."""
return features * math_ops.sigmoid(features)
class AutoMixedPrecisionTest(test.TestCase, parameterized.TestCase):
"""Tests the Grappler auto mixed precision optimizer."""
IGNORE_PERF_VAR = 'TF_AUTO_MIXED_PRECISION_GRAPH_REWRITE_IGNORE_PERFORMANCE'
# TODO(benbarsdell): Add tests for eager mode with a tf.function.
def setUp(self):
super(AutoMixedPrecisionTest, self).setUp()
# Enable the CUDA tests to be run on pre-Volta GPUs by telling the grappler
# pass to ignore performance and always transform the graph.
self._original_ignore_perf_value = os.getenv(self.IGNORE_PERF_VAR)
os.environ[self.IGNORE_PERF_VAR] = '1'
def tearDown(self):
if self._original_ignore_perf_value is not None:
os.environ[self.IGNORE_PERF_VAR] = self._original_ignore_perf_value
else:
del os.environ[self.IGNORE_PERF_VAR]
super(AutoMixedPrecisionTest, self).tearDown()
def _lower_precision_dtype(self, mode):
return dtypes.float16 if mode == 'cuda' else dtypes.bfloat16
def _assert_output_f16(self, mode, node_map, node_name, output_port=0):
self.assertEqual(node_map[node_name].output_info[output_port].dtype,
self._lower_precision_dtype(mode).as_datatype_enum)
def _run(self, mode, fetches):
"""Runs the graph and returns the evaluation of the fetches."""
with session.Session(config=_get_config(None)) as sess:
sess.run(variables.global_variables_initializer())
output_val_ref = self.evaluate(fetches)
with session.Session(config=_get_config(mode)) as sess:
sess.run(variables.global_variables_initializer())
metadata = config_pb2.RunMetadata()
output_val = sess.run(fetches, run_metadata=metadata)
return output_val_ref, output_val, metadata.cost_graph
def _maybe_skip(self, mode):
if mode == 'cuda' and not test.is_gpu_available(cuda_only=True):
self.skipTest('No GPU is available')
if mode == 'mkl' and not test_util.IsMklEnabled():
self.skipTest('MKL is not enabled')
# Test will fail on machines without AVX512f, e.g., Broadwell
is_avx512f = _pywrap_utils.IsDataTypeSupportedByOneDNNOnThisCPU(
dtypes.bfloat16
)
if mode == 'mkl' and not is_avx512f:
self.skipTest('Skipping test due to non-AVX512f machine')
def _run_simple_loop_test(self, mode, inp, body, out):
"""Runs a test of a simple loop.
The loop has different node colors in different sections of the graph. The
arguments must be strings where each character represents the color of a
node in that section of the graph: w = allow, g = infer, c = clear,
b = deny. CAPITALIZED characters indicate that the node is expected to be
changed to DT_HALF during graph optimization.
inp -> loop [ body ] -> out.
Args:
mode: Either 'cuda' or 'mkl'.
inp: A string of letters indicating the colors and expected dtypes of the
input nodes.
body: A string of letters indicating the colors and expected dtypes of the
body nodes.
out: A string of letters indicating the colors and expected dtypes of the
output nodes.
"""
self._maybe_skip(mode)
with ops.device(_get_device(mode)):
random_seed.set_random_seed(0)
expected_types = []
for section in [inp, body, out]:
section_expected_types = []
for color in section:
if color.isupper():
expected_type = self._lower_precision_dtype(mode).as_datatype_enum
else:
expected_type = types_pb2.DT_FLOAT
section_expected_types.append(expected_type)
expected_types.append(section_expected_types)
a = _build_simple_loop_graph(inp, body, out)
output_val_ref, output_val, cost_graph = self._run(mode, a)
node_map = _build_node_map(cost_graph.node)
section_names = ['input', 'while/body', 'output']
all_types_correct = True
for section_name, expected_types in zip(section_names, expected_types):
for i, expected_type in enumerate(expected_types):
node_name = section_name + '_%i' % i
output_port = 0
optimized_type = node_map[node_name].output_info[output_port].dtype
if optimized_type != expected_type:
print('Expected node %s to have type %s but got type %s' %
(node_name, expected_type, optimized_type))
all_types_correct = False
self.assertTrue(all_types_correct)
if mode == 'mkl':
self.assertAllClose(output_val_ref, output_val, atol=2e-2, rtol=2e-2)
else:
self.assertAllClose(output_val_ref, output_val, atol=2e-3, rtol=1e-3)
@parameterized.parameters(['cuda', 'mkl'])
@test_util.run_deprecated_v1
@test_util.disable_xla('This test does not pass with XLA')
def test_conv_bn(self, mode):
"""Test graph with convolution followed by batch norm."""
self._maybe_skip(mode)
with ops.device(_get_device(mode)):
random_seed.set_random_seed(0)
x = _input([2, 8, 8, 1])
x = _conv_bn(x)
output = _conv_bn(x)
output_val_ref, output_val, cost_graph = self._run(mode, output)
node_map = _build_node_map(cost_graph.node)
num_to_f16, num_to_fp32 = _count_casts(mode, cost_graph.node)
self._assert_output_f16(mode, node_map, 'Conv2D')
self._assert_output_f16(mode, node_map, 'FusedBatchNormV3')
self._assert_output_f16(mode, node_map, 'Conv2D_1')
self.assertEqual(num_to_f16, 3) # Before Conv2D:0, Conv2D:1, Conv2D_1:1
self.assertEqual(num_to_fp32, 1) # After FusedBatchNormV3:0
if mode == 'mkl':
tol = 1e-2
elif test.is_built_with_rocm():
# Bump up the tolerance for the ROCm platform
# The default tolerance (1e-3) results in a tiny fraction (<1%) of
# miscompares on ROCm platform, and hence the tolerance bump
tol = 2e-3
else:
tol = 1e-3
self.assertAllClose(output_val_ref, output_val, atol=tol, rtol=tol)
@parameterized.parameters(['cuda', 'mkl'])
@test_util.run_deprecated_v1
@test_util.disable_xla('This test does not pass with XLA')
def test_conv3d_bn(self, mode):
"""Test graph with convolution followed by batch norm."""
self._maybe_skip(mode)
if mode == 'cuda':
# TODO(reedwm): enable these tests when cuDNN is upgraded to >= 7.6.2.
self.skipTest('Test case should be skipped when cuDNN < 7.6.2')
with ops.device(_get_device(mode)):
random_seed.set_random_seed(0)
x = _input([2, 8, 8, 8, 1])
x = _conv3d_bn(x)
output = _conv3d_bn(x)
output_val_ref, output_val, cost_graph = self._run(mode, output)
node_map = _build_node_map(cost_graph.node)
num_to_fp16, num_to_fp32 = _count_casts(mode, cost_graph.node)
self._assert_output_f16(mode, node_map, 'Conv3D')
self._assert_output_f16(mode, node_map, 'FusedBatchNormV3')
self._assert_output_f16(mode, node_map, 'Conv3D_1')
self.assertEqual(num_to_fp16, 3) # Before Conv3D:0, Conv3D:1, Conv3D_1:1
self.assertEqual(num_to_fp32, 1) # After FusedBatchNormV3:0
self.assertAllClose(output_val_ref, output_val, atol=1e-2, rtol=1e-2)
@parameterized.parameters(['cuda', 'mkl'])
@test_util.run_deprecated_v1
@test_util.disable_xla('This test does not pass with XLA')
def test_conv3d(self, mode):
"""Test grad ops with convolution3d graph."""
self._maybe_skip(mode)
if mode == 'cuda':
# TODO(reedwm): enable these tests when cuDNN is upgraded to >= 7.6.2.
self.skipTest('Test case should be skipped when cuDNN < 7.6.2')
with ops.device(_get_device(mode)):
random_seed.set_random_seed(0)
x = _input([2, 8, 8, 8, 1])
f = _weight([3, 3, 3, 1, 6])
y = _conv3d(x, f)
y = array_ops.identity(y)
optimizer = gradient_descent.GradientDescentOptimizer(learning_rate=0.01)
g = optimizer.compute_gradients(y, [x, f])
output = (y, g)
output_val_ref, output_val, cost_graph = self._run(mode, output)
node_map = _build_node_map(cost_graph.node)
self._assert_output_f16(mode, node_map, 'Conv3D')
self._assert_output_f16(mode, node_map,
'gradients/Conv3D_grad/Conv3DBackpropInputV2')
self._assert_output_f16(mode, node_map,
'gradients/Conv3D_grad/Conv3DBackpropFilterV2')
output_val_ref, output_val, cost_graph = self._run(mode, output)
tol = 5e-2 if mode == 'mkl' else 1e-3
self.assertAllClose(output_val_ref, output_val, atol=tol, rtol=tol)
@parameterized.parameters(['cuda', 'mkl'])
@test_util.run_deprecated_v1
@test_util.disable_xla('This test does not pass with XLA')
def test_conv_bn_dropout(self, mode):
"""Test dropout precision of convolution batch norm graph."""
self._maybe_skip(mode)
with ops.device(_get_device(mode)):
random_seed.set_random_seed(0)
x = _input([2, 8, 8, 1])
y = _conv_bn(x)
y = nn.dropout(y, rate=0.5)
y = math_ops.add(y, 1, name='addition')
y = _conv_bn(y)
y = array_ops.identity(y)
optimizer = gradient_descent.GradientDescentOptimizer(learning_rate=0.01)
g = optimizer.compute_gradients(y, [x])
output = (y, g)
output_val_ref, output_val, cost_graph = self._run(mode, output)
node_map = _build_node_map(cost_graph.node)
self._assert_output_f16(mode, node_map, 'Conv2D')
self._assert_output_f16(mode, node_map, 'FusedBatchNormV3')
# We do not assert dropout's dtype because we do not want to rely on the
# node names of dropout's internal implementation.
self._assert_output_f16(mode, node_map, 'addition')
self._assert_output_f16(mode, node_map, 'Conv2D_1')
output_val_ref, output_val, cost_graph = self._run(mode, output)
# Bump up the tolerance for the ROCm platform
# The default tolerance (1e-3) results in a tiny fraction (<1%) of
# miscompares on ROCm platform, and hence the tolerance bump
tol = 2e-3 if test.is_built_with_rocm else 1e-3
tol = 5e-2 if mode == 'mkl' else tol
self.assertAllClose(output_val_ref, output_val, atol=tol, rtol=tol)
# TODO(reedwm): Fix and enable this test with MKL. Currently this crashes with
# MKL
@parameterized.parameters(['cuda'])
@test_util.run_deprecated_v1
@test_util.disable_xla('This test does not pass with XLA')
def test_conv_pool(self, mode):
"""Test graph with convolution followed by pooling."""
self._maybe_skip(mode)
with ops.device(_get_device(mode)):
random_seed.set_random_seed(0)
x = _input([2, 8, 8, 1])
output = _conv_pool(x)
output_val_ref, output_val, cost_graph = self._run(mode, output)
node_map = _build_node_map(cost_graph.node)
num_to_f16, num_to_fp32 = _count_casts(mode, cost_graph.node)
self._assert_output_f16(mode, node_map, 'Conv2D')
self._assert_output_f16(mode, node_map, 'Relu')
self._assert_output_f16(mode, node_map, 'MaxPool')
self._assert_output_f16(mode, node_map, 'Conv2D_1')
self.assertEqual(num_to_f16, 5)
self.assertEqual(num_to_fp32, 1)
tol = 5e-3 if mode == 'mkl' else 1e-3
self.assertAllClose(output_val_ref, output_val, atol=tol, rtol=tol)
# TODO(benbarsdell): This test has not been tried with MKL.
@parameterized.parameters(['cuda'])
@test_util.run_deprecated_v1
@test_util.disable_xla('This test does not pass with XLA')
def test_depthwise_conv2d(self, mode):
"""Test grad ops with depthwise convolution2d graph."""
self._maybe_skip(mode)
cudnn_version_str = sysconfig_lib.get_build_info().get(
'cudnn_version', '0.0')
cudnn_version = tuple([int(x) for x in cudnn_version_str.split('.')])
if cudnn_version < (8,):
# Depthwise conv2d ops are only enabled in auto_mixed_precision as of
# cuDNN v8.
self.skipTest('cuDNN version >= 8 required')
with ops.device(_get_device(mode)):
random_seed.set_random_seed(0)
x = _input([2, 8, 8, 1])
f = _weight([3, 3, 1, 4])
y = _depthwise_conv2d(x, f)
y = array_ops.identity(y)
optimizer = gradient_descent.GradientDescentOptimizer(learning_rate=0.01)
g = optimizer.compute_gradients(y, [x, f])
output = (y, g)
output_val_ref, output_val, cost_graph = self._run(mode, output)
node_map = _build_node_map(cost_graph.node)
self._assert_output_f16(mode, node_map, 'depthwise')
self._assert_output_f16(
mode, node_map,
'gradients/depthwise_grad/DepthwiseConv2dNativeBackpropInput')
self._assert_output_f16(
mode, node_map,
'gradients/depthwise_grad/DepthwiseConv2dNativeBackpropFilter')
output_val_ref, output_val, cost_graph = self._run(mode, output)
tol = 2e-3
self.assertAllClose(output_val_ref, output_val, atol=tol, rtol=tol)
@parameterized.parameters(['cuda', 'mkl'])
@test_util.run_v1_only('b/138749235')
@test_util.disable_xla('This test does not pass with XLA')
def test_simple_loop(self, mode):
"""Test graph with while loop."""
self._maybe_skip(mode)
with ops.device(_get_device(mode)):
random_seed.set_random_seed(0)
x = _input([8, 8])
y = _simple_loop(x, _matmul_act)[1]
optimizer = gradient_descent.GradientDescentOptimizer(learning_rate=0.01)
g = optimizer.compute_gradients(y, [x])
output = (y, g)
output_val_ref, output_val, cost_graph = self._run(mode, output)
node_map = _build_node_map(cost_graph.node)
self._assert_output_f16(mode, node_map, 'while/MatMul')
self._assert_output_f16(mode, node_map, 'while/Relu')
tol = 1e-2 if mode == 'mkl' else 1e-3
self.assertAllClose(output_val_ref, output_val, atol=tol, rtol=tol)
@parameterized.parameters(['cuda', 'mkl'])
@test_util.run_v1_only('b/138749235')
@test_util.disable_xla('This test does not pass with XLA')
def test_loop_with_vars_intertwined(self, mode):
"""Test graph with intertwined while loops."""
self._maybe_skip(mode)
with ops.device(_get_device(mode)):
random_seed.set_random_seed(0)
x = _input([8, 8])
_, _, k, l = _loop_vars_intertwined(
array_ops.ones(array_ops.shape(x)), x, _matmul_act, _matmul_act)
optimizer = gradient_descent.GradientDescentOptimizer(learning_rate=0.01)
g = optimizer.compute_gradients(k, [x])
output = (k, l, g)
output_val_ref, output_val, cost_graph = self._run(mode, output)
node_map = _build_node_map(cost_graph.node)
self._assert_output_f16(mode, node_map, 'while/MatMul')
self._assert_output_f16(mode, node_map, 'while/Relu')
self._assert_output_f16(mode, node_map, 'while/MatMul_1')
self._assert_output_f16(mode, node_map, 'while/Relu_1')
tol = 5e-3 if mode == 'mkl' else 1e-3
self.assertAllClose(output_val_ref, output_val, atol=tol, rtol=tol)
@parameterized.parameters(['cuda'])
@test_util.run_deprecated_v1
@test_util.disable_xla('This test does not pass with XLA')
def test_multi_paths(self, mode):
"""Test graph with multiple paths."""
self._maybe_skip(mode)
with ops.device(_get_device(mode)):
random_seed.set_random_seed(0)
x = _input([2, 8, 8, 3])
x1, x2, x3 = array_ops.split(x, num_or_size_splits=3, axis=3)
y1 = _conv_pool(x1)
y2 = _conv_pool(x2)
y3 = _conv_pool(x3)
y = array_ops.concat([y1, y2, y3], axis=3)
y = array_ops.identity(y)
optimizer = gradient_descent.GradientDescentOptimizer(learning_rate=0.01)
g = optimizer.compute_gradients(y, [x])
output = (y, g)
output_val_ref, output_val, cost_graph = self._run(mode, output)
node_map = _build_node_map(cost_graph.node)
self._assert_output_f16(mode, node_map, 'split')
for suffix in [''] + ['_%i' % i for i in range(1, 6)]:
self._assert_output_f16(mode, node_map, 'Conv2D' + suffix)
self._assert_output_f16(mode, node_map, 'Relu' + suffix)
self._assert_output_f16(mode, node_map, 'MaxPool' + suffix)
self._assert_output_f16(mode, node_map, 'concat')
atol = 1e-2 if test.is_built_with_rocm() else 1e-3
self.assertAllClose(output_val_ref, output_val, atol=atol, rtol=1e-3)
@parameterized.parameters(['cuda', 'mkl'])
@test_util.run_deprecated_v1
@test_util.disable_xla('This test does not pass with XLA')
def test_multi_paths_2(self, mode):
"""Test graph with multiple paths."""
self._maybe_skip(mode)
with ops.device(_get_device(mode)):
random_seed.set_random_seed(0)
x = _input([8, 8])
y1 = _matmul_act(x)
y2 = _matmul_act(x)
y = y1 + y2 + x
optimizer = gradient_descent.GradientDescentOptimizer(learning_rate=0.01)
g = optimizer.compute_gradients(y, [x])
output = (g, y)
output_val_ref, output_val, cost_graph = self._run(mode, output)
node_map = _build_node_map(cost_graph.node)
self._assert_output_f16(mode, node_map, 'MatMul')
self._assert_output_f16(mode, node_map, 'Relu')
self._assert_output_f16(mode, node_map, 'MatMul_1')
self._assert_output_f16(mode, node_map, 'Relu_1')
if mode == 'mkl':
tol = 2e-2
elif test.is_built_with_rocm():
# Bump up the tolerance for the ROCm platform
# The default tolerance (1e-3) results in a tiny fraction (<1%) of
# miscompares on ROCm platform, and hence the tolerance bump
tol = 1e-2
else:
tol = 1e-3
self.assertAllClose(output_val_ref, output_val, atol=tol, rtol=tol)
@parameterized.parameters(['cuda']) # MKL doesn't support bf16 Sigmoid
@test_util.run_v1_only('b/138749235')
@test_util.disable_xla('This test does not pass with XLA')
def test_recurrent_lstm(self, mode):
"""Test graph with recurrent lstm."""
self._maybe_skip(mode)
with ops.device(_get_device(mode)):
random_seed.set_random_seed(0)
init_c = _input([8, 4])
init_h = _input([8, 4])
_, _, h, _ = _recurrent_lstm(init_c, init_h)
optimizer = gradient_descent.GradientDescentOptimizer(learning_rate=0.01)
g = optimizer.compute_gradients(h, [init_c, init_h])
output = (h, g)
output_val_ref, output_val, cost_graph = self._run(mode, output)
node_map = _build_node_map(cost_graph.node)
self._assert_output_f16(mode, node_map, 'while/concat')
self._assert_output_f16(mode, node_map, 'while/MatMul')
self._assert_output_f16(mode, node_map, 'while/split')
self._assert_output_f16(mode, node_map, 'while/Sigmoid')
self._assert_output_f16(mode, node_map, 'while/Sigmoid_1')
self._assert_output_f16(mode, node_map, 'while/Sigmoid_2')
self._assert_output_f16(mode, node_map, 'while/Tanh')
self._assert_output_f16(mode, node_map, 'while/Tanh_1')
self.assertAllClose(output_val_ref, output_val, atol=1e-3, rtol=1e-3)
@parameterized.parameters(['cuda', 'mkl'])
@test_util.run_v1_only('v1 loop test')
@test_util.disable_xla('This test does not pass with XLA')
def test_propagation_through_simple_loop_1(self, mode):
self._run_simple_loop_test(mode, 'W', 'C', 'C')
@parameterized.parameters(['cuda', 'mkl'])
@test_util.run_v1_only('v1 loop test')
@test_util.disable_xla('This test does not pass with XLA')
def test_propagation_through_simple_loop_2(self, mode):
self._run_simple_loop_test(mode, 'C', 'C', 'W')
@parameterized.parameters(['cuda', 'mkl'])
@test_util.run_v1_only('v1 loop test')
@test_util.disable_xla('This test does not pass with XLA')
def test_propagation_through_simple_loop_3(self, mode):
self._run_simple_loop_test(mode, 'W', 'G', 'W')
@parameterized.parameters(['cuda', 'mkl'])
@test_util.run_v1_only('v1 loop test')
@test_util.disable_xla('This test does not pass with XLA')
def test_propagation_through_simple_loop_4(self, mode):
self._run_simple_loop_test(mode, 'W', 'gbg', 'W')
@parameterized.parameters(['cuda', 'mkl'])
@test_util.run_v1_only('b/138749235')
@test_util.disable_xla('This test does not pass with XLA')
def test_propagation_through_simple_loop_5(self, mode):
self._run_simple_loop_test(mode, 'b', 'gWC', 'c')
@parameterized.parameters(['cuda', 'mkl'])
@test_util.run_v1_only('b/138749235')
@test_util.disable_xla('This test does not pass with XLA')
def test_propagation_through_simple_loop_6(self, mode):
self._run_simple_loop_test(mode, 'b', 'CWCG', 'C')
@parameterized.parameters(['cuda', 'mkl'])
@test_util.run_v1_only('b/138749235')
@test_util.disable_xla('This test does not pass with XLA')
def test_propagation_through_simple_loop_7(self, mode):
self._run_simple_loop_test(mode, 'C', 'GWCG', 'C')
@parameterized.parameters(['cuda', 'mkl'])
@test_util.run_v1_only('b/138749235')
@test_util.disable_xla('This test does not pass with XLA')
def test_propagation_through_simple_loop_8(self, mode):
self._run_simple_loop_test(mode, 'C', 'CgbgWC', 'g')
@parameterized.parameters(['cuda', 'mkl'])
@test_util.run_deprecated_v1
@test_util.disable_xla('This test does not pass with XLA')
def test_noninlined_funcdef(self, mode):
"""Test graph with non-inlined function subgraph.
This requires the grappler pass to handle an OpDef that only appears in the
graph's function registry instead of the global op registry.
Args:
mode: Either 'cuda' or 'mkl'.
"""
self._maybe_skip(mode)
with ops.device(_get_device(mode)):
random_seed.set_random_seed(0)
x = _input([8, 8])
y = _matmul_act(x)
y = _example_noninlined_funcdef(y)
optimizer = gradient_descent.GradientDescentOptimizer(learning_rate=0.01)
g = optimizer.compute_gradients(y, [x])
output = (g, y)
output_val_ref, output_val, cost_graph = self._run(mode, output)
node_map = _build_node_map(cost_graph.node)
self._assert_output_f16(mode, node_map, 'MatMul')
tol = 1e-2 if mode == 'mkl' else 1e-3
atol = 1e-2 if test.is_built_with_rocm() else tol
self.assertAllClose(output_val_ref, output_val, atol=atol, rtol=tol)
@parameterized.parameters(['cuda', 'mkl'])
@test_util.run_deprecated_v1
@test_util.disable_xla('This test does not pass with XLA')
def test_ingraph_train_loop(self, mode):
"""Tests a graph containing a while loop around a training update.
This requires the grappler pass to take special care with its handling of
Enter ops that appear in front of reads from non-resource variables. See
the use of NodeImplicitlyReadsVariable in auto_mixed_precision.cc.
Args:
mode: Either 'cuda' or 'mkl'.
"""
self._maybe_skip(mode)
if tf2.enabled():
# This test tests non-resource variables, which are only used in TF1.
self.skipTest('TensorFlow 1 required')
with ops.device(_get_device(mode)):
random_seed.set_random_seed(1234)
np.random.seed(1234)
num_iter, bs, nchan, nclass = 100, 64, 32, 100
data = np.random.normal(size=(bs * num_iter, nchan)).astype(np.float32)
labels = np.random.randint(nclass, size=(bs * num_iter,))
ds = dataset_ops.Dataset.from_tensor_slices((data, labels))
ds = ds.batch(bs).prefetch(3)
it = ds.make_one_shot_iterator()
def body(_, i):
i += 1
x, yt = it.get_next()
dense = layers.Dense(nclass)
y = dense(x)
loss = losses.sparse_softmax_cross_entropy(yt, y)
opt = adam.AdamOptimizer()
train_op = opt.minimize(loss, var_list=dense.trainable_weights)
with ops.control_dependencies([train_op]):
loss = array_ops.identity(loss)
return loss, i
begin, end = constant_op.constant(0), constant_op.constant(num_iter)
loss, _ = while_loop.while_loop(lambda loss, i: math_ops.less(i, end),
body, [0.0, begin])
output_val_ref, output_val, cost_graph = self._run(mode, loss)
node_map = _build_node_map(cost_graph.node)
self._assert_output_f16(mode, node_map, 'while/dense/MatMul')
self._assert_output_f16(mode, node_map,
'while/gradients/while/dense/MatMul_grad/MatMul_1')
self.assertAllClose(output_val_ref, output_val, atol=1e-3, rtol=1e-3)
# TODO(benbarsdell): Add tests for list ops (TensorList*) that pass through
# graph source/sink nodes, similar to the TensorListThroughFunction C++ test.
# Tests here will have the advantage of catching changes in the types of ops
# that are added to the graph.
if __name__ == '__main__':
test.main()
tensorflow/python/grappler/cluster.py
+118
View File
@@ -0,0 +1,118 @@
# 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.
# ==============================================================================
"""A python interface for Grappler clusters."""
import contextlib
from tensorflow.core.framework import step_stats_pb2
from tensorflow.core.grappler.costs import op_performance_data_pb2
from tensorflow.core.protobuf import device_properties_pb2
from tensorflow.python.grappler import _pywrap_tf_cluster as tf_cluster
class Cluster(object):
"""Grappler Clusters."""
def __init__(self,
allow_soft_placement=True,
disable_detailed_stats=True,
disable_timeline=True,
devices=None):
"""Creates a Cluster.
Args:
allow_soft_placement: If True, TF will automatically fix illegal
placements instead of erroring out if the placement isn't legal.
disable_detailed_stats: If True, detailed statistics will not be
available.
disable_timeline: If True, the timeline information will not be reported.
devices: A list of devices of type device_properties_pb2.NamedDevice.
If None, a device list will be created based on the spec of
the local machine.
"""
self._tf_cluster = None
self._generate_timeline = not disable_timeline
if devices is None:
self._tf_cluster = tf_cluster.TF_NewCluster(allow_soft_placement,
disable_detailed_stats)
else:
devices_serialized = [device.SerializeToString() for device in devices]
self._tf_cluster = tf_cluster.TF_NewVirtualCluster(devices_serialized)
def Shutdown(self):
if self._tf_cluster is not None:
tf_cluster.TF_ShutdownCluster(self._tf_cluster)
self._tf_cluster = None
def __del__(self):
self.Shutdown()
@property
def tf_cluster(self):
return self._tf_cluster
def ListDevices(self):
"""Returns a list of available hardware devices."""
if self._tf_cluster is None:
return []
return [device_properties_pb2.NamedDevice.FromString(device)
for device in tf_cluster.TF_ListDevices(self._tf_cluster)]
def ListAvailableOps(self):
"""Returns a list of all available operations (sorted alphabetically)."""
return tf_cluster.TF_ListAvailableOps()
def GetSupportedDevices(self, item):
return tf_cluster.TF_GetSupportedDevices(self._tf_cluster, item.tf_item)
def EstimatePerformance(self, device):
return tf_cluster.TF_EstimatePerformance(device.SerializeToString())
def MeasureCosts(self, item):
"""Returns the cost of running the specified item.
Args:
item: The item for which to measure the costs.
Returns: The triplet op_perfs, runtime, step_stats.
"""
op_perf_bytes_list, run_time, step_stats_bytes = tf_cluster.TF_MeasureCosts(
item.tf_item, self._tf_cluster, self._generate_timeline)
op_perfs = [op_performance_data_pb2.OpPerformance.FromString(op_perf_bytes)
for op_perf_bytes in op_perf_bytes_list]
return (op_perfs, run_time,
step_stats_pb2.StepStats.FromString(step_stats_bytes))
def DeterminePeakMemoryUsage(self, item):
"""Returns a snapshot of the peak memory usage.
Args:
item: The item for which to measure the costs.
Returns: A hashtable indexed by device name.
"""
return tf_cluster.TF_DeterminePeakMemoryUsage(item.tf_item,
self._tf_cluster)
@contextlib.contextmanager
def Provision(allow_soft_placement=True,
disable_detailed_stats=True,
disable_timeline=True,
devices=None):
cluster = Cluster(allow_soft_placement, disable_detailed_stats,
disable_timeline, devices)
yield cluster
cluster.Shutdown()
tensorflow/python/grappler/cluster_test.py
+182
View File
@@ -0,0 +1,182 @@
# 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 the swig wrapper of clusters."""
from tensorflow.core.protobuf import device_properties_pb2
from tensorflow.python.framework import meta_graph
from tensorflow.python.framework import ops
from tensorflow.python.grappler import cluster
from tensorflow.python.grappler import item
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import random_ops
from tensorflow.python.platform import test
class ClusterTest(test.TestCase):
def testBasic(self):
with ops.Graph().as_default() as g:
a = random_ops.random_uniform(shape=())
b = random_ops.random_uniform(shape=())
c = a + b
train_op = ops.get_collection_ref(ops.GraphKeys.TRAIN_OP)
train_op.append(c)
mg = meta_graph.create_meta_graph_def(graph=g)
grappler_item = item.Item(mg)
grappler_cluster = cluster.Cluster(
disable_detailed_stats=False, disable_timeline=False)
op_perfs, run_time, step_stats = grappler_cluster.MeasureCosts(
grappler_item)
self.assertTrue(run_time > 0)
self.assertEqual(len(op_perfs), 4)
self.assertTrue(step_stats.dev_stats)
def testNoDetailedStats(self):
with ops.Graph().as_default() as g:
a = random_ops.random_uniform(shape=())
b = random_ops.random_uniform(shape=())
c = a + b
train_op = ops.get_collection_ref(ops.GraphKeys.TRAIN_OP)
train_op.append(c)
mg = meta_graph.create_meta_graph_def(graph=g)
grappler_item = item.Item(mg)
grappler_cluster = cluster.Cluster(disable_detailed_stats=True)
op_perfs, run_time, step_stats = grappler_cluster.MeasureCosts(
grappler_item)
self.assertTrue(run_time > 0)
self.assertEqual(len(op_perfs), 0)
self.assertEqual(len(step_stats.dev_stats), 0)
def testMemoryEstimates(self):
with ops.Graph().as_default() as g:
with ops.device('/job:localhost/replica:0/task:0/device:CPU:0'):
a = random_ops.random_uniform(shape=())
b = random_ops.random_uniform(shape=())
c = a + b
train_op = ops.get_collection_ref(ops.GraphKeys.TRAIN_OP)
train_op.append(c)
mg = meta_graph.create_meta_graph_def(graph=g)
grappler_item = item.Item(mg)
grappler_cluster = cluster.Cluster(
disable_detailed_stats=True, disable_timeline=True)
peak_mem = grappler_cluster.DeterminePeakMemoryUsage(grappler_item)
self.assertLessEqual(1, len(peak_mem))
snapshot = peak_mem['/job:localhost/replica:0/task:0/device:CPU:0']
peak_usage = snapshot[0]
self.assertEqual(12, peak_usage)
live_tensors = snapshot[1]
self.assertEqual(5, len(live_tensors))
def testVirtualCluster(self):
with ops.Graph().as_default() as g:
with ops.device('/device:GPU:0'):
a = random_ops.random_uniform(shape=[1024, 1024])
b = random_ops.random_uniform(shape=[1024, 1024])
c = a + b
train_op = ops.get_collection_ref(ops.GraphKeys.TRAIN_OP)
train_op.append(c)
mg = meta_graph.create_meta_graph_def(graph=g)
grappler_item = item.Item(mg)
device_properties = device_properties_pb2.DeviceProperties(
type='GPU',
frequency=1000,
num_cores=60,
environment={'architecture': '7'})
named_device = device_properties_pb2.NamedDevice(
properties=device_properties, name='/device:GPU:0')
grappler_cluster = cluster.Cluster(
disable_detailed_stats=False,
disable_timeline=False,
devices=[named_device])
op_perfs, run_time, _ = grappler_cluster.MeasureCosts(grappler_item)
self.assertEqual(run_time, 0.000209)
self.assertEqual(len(op_perfs), 5)
estimated_perf = grappler_cluster.EstimatePerformance(named_device)
self.assertEqual(7680.0, estimated_perf)
def testContext(self):
with ops.Graph().as_default() as g:
a = random_ops.random_uniform(shape=())
b = random_ops.random_uniform(shape=())
c = a + b
train_op = ops.get_collection_ref(ops.GraphKeys.TRAIN_OP)
train_op.append(c)
mg = meta_graph.create_meta_graph_def(graph=g)
grappler_item = item.Item(mg)
with cluster.Provision(
disable_detailed_stats=False, disable_timeline=False) as gcluster:
op_perfs, run_time, step_stats = gcluster.MeasureCosts(grappler_item)
self.assertTrue(run_time > 0)
self.assertEqual(len(op_perfs), 4)
self.assertTrue(step_stats.dev_stats)
def testAvailableOps(self):
with cluster.Provision() as gcluster:
op_names = gcluster.ListAvailableOps()
self.assertTrue('Add' in op_names)
self.assertTrue('MatMul' in op_names)
self.assertEqual(op_names, sorted(op_names))
def testSupportDevices(self):
with ops.Graph().as_default() as g:
a = random_ops.random_uniform(shape=(2, 3))
b = random_ops.random_uniform(shape=(2, 3))
c = a + b
dims = math_ops.range(0, array_ops.rank(c), 1)
d = math_ops.reduce_sum(a, axis=dims)
train_op = ops.get_collection_ref(ops.GraphKeys.TRAIN_OP)
train_op.append(d)
mg = meta_graph.create_meta_graph_def(graph=g)
grappler_item = item.Item(mg)
device_properties = device_properties_pb2.DeviceProperties(
type='GPU', frequency=1000, num_cores=60)
named_gpu = device_properties_pb2.NamedDevice(
properties=device_properties, name='/GPU:0')
device_properties = device_properties_pb2.DeviceProperties(
type='CPU', frequency=3000, num_cores=6)
named_cpu = device_properties_pb2.NamedDevice(
properties=device_properties, name='/CPU:0')
virtual_cluster = cluster.Cluster(devices=[named_cpu, named_gpu])
supported_dev = virtual_cluster.GetSupportedDevices(grappler_item)
self.assertEqual(supported_dev['add'], ['/CPU:0', '/GPU:0'])
self.assertEqual(supported_dev['Sum'], ['/CPU:0', '/GPU:0'])
self.assertEqual(supported_dev['range'], ['/CPU:0', '/GPU:0'])
real_cluster = cluster.Cluster()
supported_dev = real_cluster.GetSupportedDevices(grappler_item)
if test.is_gpu_available():
self.assertEqual(supported_dev['add'], [
'/job:localhost/replica:0/task:0/device:CPU:0',
'/job:localhost/replica:0/task:0/device:GPU:0'
])
self.assertEqual(supported_dev['Sum'], [
'/job:localhost/replica:0/task:0/device:CPU:0',
'/job:localhost/replica:0/task:0/device:GPU:0'
])
# The axis tensor must reside on the host
self.assertEqual(supported_dev['range'],
['/job:localhost/replica:0/task:0/device:CPU:0'])
else:
self.assertEqual(supported_dev['add'],
['/job:localhost/replica:0/task:0/device:CPU:0'])
if __name__ == '__main__':
test.main()
tensorflow/python/grappler/cluster_wrapper.cc
+335
View File
@@ -0,0 +1,335 @@
/* 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 <algorithm>
#include <cfloat>
#include <cstddef>
#include <cstdint>
#include <iterator>
#include <memory>
#include <set>
#include <stdexcept>
#include <string>
#include <tuple>
#include <unordered_map>
#include <vector>
#include "absl/log/check.h"
#include "absl/status/status.h"
#include "pybind11/pybind11.h" // from @pybind11
#include "pybind11/stl.h" // from @pybind11
#include "tensorflow/core/framework/kernel_def.pb.h"
#include "tensorflow/core/framework/memory_types.h"
#include "tensorflow/core/framework/op_def.pb.h"
#include "tensorflow/core/framework/step_stats.pb.h"
#include "tensorflow/core/grappler/clusters/cluster.h"
#include "tensorflow/core/grappler/clusters/single_machine.h"
#include "tensorflow/core/grappler/clusters/virtual_cluster.h"
#include "tensorflow/core/grappler/costs/cost_estimator.h"
#include "tensorflow/core/grappler/costs/graph_memory.h"
#include "tensorflow/core/grappler/costs/measuring_cost_estimator.h"
#include "tensorflow/core/grappler/costs/op_level_cost_estimator.h"
#include "tensorflow/core/grappler/costs/op_performance_data.pb.h"
#include "tensorflow/core/grappler/costs/utils.h"
#include "tensorflow/core/grappler/devices.h"
#include "tensorflow/core/grappler/grappler_item.h"
#include "tensorflow/core/grappler/utils.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/protobuf/config.pb.h"
#include "tensorflow/core/protobuf/device_properties.pb.h"
#include "tensorflow/python/lib/core/pybind11_status.h"
namespace py = pybind11;
absl::Status _GetOpPerformanceDataAndRunTime(
const tensorflow::grappler::GrapplerItem& item,
tensorflow::grappler::CostEstimator* cost_measure,
tensorflow::OpPerformanceList* op_performance_data,
tensorflow::grappler::Costs* costs) {
absl::Status status = cost_measure->Initialize(item);
if (!status.ok()) return status;
tensorflow::RunMetadata run_metadata;
tsl::MaybeRaiseRegisteredFromStatus(
cost_measure->PredictCosts(item.graph, &run_metadata, costs));
if (op_performance_data) {
*op_performance_data = tensorflow::grappler::CostGraphToOpPerformanceData(
run_metadata.cost_graph(), item.graph);
}
return absl::OkStatus();
}
PYBIND11_MAKE_OPAQUE(tensorflow::grappler::Cluster);
PYBIND11_MODULE(_pywrap_tf_cluster, m) {
py::class_<tensorflow::grappler::Cluster> grappler_cluster(m, "Cluster");
m.def("TF_NewCluster",
[](bool allow_soft_placement,
bool disable_detailed_stats) -> tensorflow::grappler::Cluster* {
// TODO(petebu): Make these named arguments with default values
// instead.
int num_cpu_cores =
tensorflow::grappler::GetNumAvailableLogicalCPUCores();
int num_gpus = tensorflow::grappler::GetNumAvailableGPUs();
int timeout_s = 60 * 10;
std::unique_ptr<tensorflow::grappler::Cluster> cluster =
std::make_unique<tensorflow::grappler::SingleMachine>(
timeout_s, num_cpu_cores, num_gpus);
cluster->DisableDetailedStats(disable_detailed_stats);
cluster->AllowSoftPlacement(allow_soft_placement);
cluster->SetNumWarmupSteps(10);
tsl::MaybeRaiseRegisteredFromStatus(cluster->Provision());
return cluster.release();
});
m.def("TF_NewVirtualCluster",
[](const std::vector<py::bytes>& serialized_named_devices)
-> tensorflow::grappler::Cluster* {
std::vector<tensorflow::NamedDevice> named_devices;
for (const auto& s : serialized_named_devices) {
tensorflow::NamedDevice named_device;
if (!named_device.ParseFromString(std::string(s))) {
throw std::invalid_argument(
"The NamedDevice could not be parsed as a valid protocol "
"buffer");
}
named_devices.push_back(named_device);
}
std::unordered_map<std::string, tensorflow::DeviceProperties> devices;
for (const auto& named_device : named_devices) {
devices[named_device.name()] = named_device.properties();
}
std::unique_ptr<tensorflow::grappler::Cluster> cluster =
std::make_unique<tensorflow::grappler::VirtualCluster>(devices);
{
// TODO(petebu): Do we need to hold the GIL here?
py::gil_scoped_acquire acquire;
tsl::MaybeRaiseRegisteredFromStatus(cluster->Provision());
}
return cluster.release();
});
m.def("TF_ShutdownCluster", [](tensorflow::grappler::Cluster* cluster) {
// TODO(petebu): Do we need to hold the GIL here?
py::gil_scoped_acquire acquire;
(void)cluster->Shutdown();
});
m.def("TF_ListDevices",
[](tensorflow::grappler::Cluster* cluster) -> std::vector<py::bytes> {
const std::unordered_map<std::string, tensorflow::DeviceProperties>&
devices = cluster->GetDevices();
std::vector<py::bytes> named_devices;
for (auto& dev : devices) {
tensorflow::NamedDevice d;
d.set_name(dev.first);
*d.mutable_properties() = dev.second;
named_devices.push_back(d.SerializeAsString());
}
return named_devices;
});
m.def("TF_ListAvailableOps", []() -> std::vector<std::string> {
tensorflow::OpRegistry* registry = tensorflow::OpRegistry::Global();
std::vector<tensorflow::OpDef> ops;
registry->GetRegisteredOps(&ops);
std::vector<std::string> op_names;
op_names.reserve(ops.size());
for (const tensorflow::OpDef& op : ops) {
op_names.push_back(op.name());
}
std::sort(op_names.begin(), op_names.end());
return op_names;
});
m.def(
"TF_GetSupportedDevices",
[](tensorflow::grappler::Cluster* cluster,
tensorflow::grappler::GrapplerItem* item)
-> std::unordered_map<std::string, std::vector<std::string>> {
if (cluster == nullptr || item == nullptr) {
tsl::MaybeRaiseRegisteredFromStatus(absl::Status(
absl::InternalError("You need both a cluster and an "
"item to get supported devices.")));
}
const std::unordered_map<std::string, tensorflow::DeviceProperties>&
devices = cluster->GetDevices();
std::unordered_map<std::string, std::vector<std::string>> device_types;
for (const auto& dev : devices) {
device_types[dev.second.type()].push_back(dev.first);
}
std::unordered_map<std::string, std::set<std::string>>
supported_device_types;
std::unordered_map<std::string, std::set<std::string>>
device_restrictions;
for (const auto& node : item->graph.node()) {
for (const auto& dev : device_types) {
const std::string& type = dev.first;
if (cluster->type() != "single_machine") {
// The actual kernel may not be linked in this binary.
supported_device_types[node.name()].insert(type);
} else {
// Check the kernel capabilities
const tensorflow::DeviceType dev_type(type);
absl::Status s =
tensorflow::FindKernelDef(dev_type, node, nullptr, nullptr);
if (s.ok()) {
supported_device_types[node.name()].insert(type);
// Check which inputs are restricted to reside on the host.
// TODO: extends this to support outputs as well
tensorflow::MemoryTypeVector inp_mtypes;
tensorflow::MemoryTypeVector out_mtypes;
absl::Status s = tensorflow::MemoryTypesForNode(
tensorflow::OpRegistry::Global(), dev_type, node,
&inp_mtypes, &out_mtypes);
if (s.ok()) {
for (size_t i = 0; i < inp_mtypes.size(); ++i) {
if (inp_mtypes[i] == tensorflow::HOST_MEMORY) {
device_restrictions[tensorflow::grappler::NodeName(
node.input(i))]
.insert("CPU");
break;
}
}
}
}
}
}
}
std::unordered_map<std::string, std::vector<std::string>> result;
for (const auto& supported_dev : supported_device_types) {
const std::string& node = supported_dev.first;
std::set<std::string> feasible;
const auto it = device_restrictions.find(node);
if (it != device_restrictions.end()) {
const std::set<std::string>& candidates = supported_dev.second;
const std::set<std::string>& valid = it->second;
std::set_intersection(candidates.begin(), candidates.end(),
valid.begin(), valid.end(),
std::inserter(feasible, feasible.begin()));
} else {
feasible = supported_dev.second;
}
std::vector<std::string> device_names;
for (const std::string& type : feasible) {
auto it = device_types.find(type);
DCHECK(it != device_types.end());
for (const std::string& name : it->second) {
device_names.push_back(name);
}
}
result[node] = device_names;
}
return result;
});
m.def("TF_EstimatePerformance", [](const py::bytes& serialized_device) {
tensorflow::NamedDevice device;
if (!device.ParseFromString(std::string(serialized_device))) {
throw std::invalid_argument(
"The NamedDevice could not be parsed as a valid protocol buffer");
}
tensorflow::grappler::OpLevelCostEstimator estimator;
tensorflow::grappler::DeviceInfo info =
estimator.GetDeviceInfo(device.properties());
return info.gigaops;
});
m.def("TF_MeasureCosts",
[](tensorflow::grappler::GrapplerItem* item,
tensorflow::grappler::Cluster* cluster, bool generate_timeline)
-> std::tuple<std::vector<py::bytes>, double, py::bytes> {
const int num_measurements = cluster->type() == "virtual" ? 1 : 10;
tensorflow::grappler::MeasuringCostEstimator cost_measure(
cluster, num_measurements, 0);
tensorflow::OpPerformanceList op_performance_data;
tensorflow::grappler::Costs costs;
absl::Status s = _GetOpPerformanceDataAndRunTime(
*item, &cost_measure, &op_performance_data, &costs);
double run_time = FLT_MAX;
if (s.ok()) {
run_time = static_cast<double>(costs.execution_time.count()) / 1e9;
}
tensorflow::StepStats step_stats;
if (generate_timeline) {
tensorflow::RunMetadata metadata;
tsl::MaybeRaiseRegisteredFromStatus(
cluster->Run(item->graph, item->feed, item->fetch, &metadata));
step_stats = metadata.step_stats();
}
std::vector<py::bytes> op_perf_objs;
op_perf_objs.resize(op_performance_data.op_performance_size());
for (int i = 0; i < op_performance_data.op_performance_size(); i++) {
op_perf_objs[i] =
op_performance_data.op_performance(i).SerializeAsString();
}
py::bytes step_stats_str = step_stats.SerializeAsString();
return std::make_tuple(op_perf_objs, run_time, step_stats_str);
});
using DurationType = tensorflow::grappler::Costs::Duration::rep;
using MemoryUsage =
std::tuple<std::string, int, size_t, DurationType, DurationType>;
m.def(
"TF_DeterminePeakMemoryUsage",
[](tensorflow::grappler::GrapplerItem* item,
tensorflow::grappler::Cluster* cluster)
-> std::unordered_map<std::string,
std::tuple<int64_t, std::vector<MemoryUsage>>> {
if (item == nullptr || cluster == nullptr) {
tsl::MaybeRaiseRegisteredFromStatus(absl::Status(absl::InternalError(
"You need both a cluster and an item to determine peak "
"memory usage.")));
}
tensorflow::grappler::GraphMemory memory(*item);
if (cluster->DetailedStatsEnabled()) {
tsl::MaybeRaiseRegisteredFromStatus(memory.InferDynamically(cluster));
} else {
tsl::MaybeRaiseRegisteredFromStatus(
memory.InferStatically(cluster->GetDevices()));
}
std::unordered_map<std::string,
std::tuple<int64_t, std::vector<MemoryUsage>>>
result;
for (const auto& device : cluster->GetDevices()) {
const tensorflow::grappler::GraphMemory::MemoryUsage& usage =
memory.GetPeakMemoryUsage(device.first);
std::vector<MemoryUsage> per_device;
for (size_t i = 0; i < usage.live_tensors.size(); ++i) {
const auto& live_tensor = usage.live_tensors[i];
per_device.push_back(std::make_tuple(
live_tensor.node, live_tensor.output_id,
live_tensor.memory_used, live_tensor.allocation_time.count(),
live_tensor.deallocation_time.count()));
}
result[device.first] = std::make_tuple(usage.used_memory, per_device);
}
return result;
});
}
tensorflow/python/grappler/constant_folding_test.py
+107
View File
@@ -0,0 +1,107 @@
# 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 Grappler Constant Folding."""
import numpy as np
from tensorflow.python.eager import backprop
from tensorflow.python.eager import context
from tensorflow.python.eager import def_function
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 functional_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.ops import while_loop
from tensorflow.python.platform import test
class ConstantFoldingTest(test.TestCase):
# See b/76008022.
def testScanInsideWhile(self):
def loop_cond(idx_step, *unused_args):
return idx_step < 1
def loop_body(idx_step, y):
x = array_ops.zeros([10, 20, 30], dtype=dtypes.float32)
x = functional_ops.scan(
math_ops.add,
x,
initializer=array_ops.zeros([20, 30], dtype=dtypes.float32),
back_prop=False,
parallel_iterations=1)
with ops.device('/cpu:0'):
y = array_ops.identity(x)
return idx_step + 1, y
if test.is_gpu_available(cuda_only=True):
init_y = array_ops.zeros([10, 20, 30], dtype=dtypes.float32)
_, y = while_loop.while_loop(
loop_cond,
loop_body,
loop_vars=[0, init_y],
back_prop=False,
parallel_iterations=1)
y_v = self.evaluate(y)
self.assertAllEqual(np.zeros([10, 20, 30]), y_v)
# See b/159753857.
def testGradientGraphOptimization(self):
@def_function.function
def f(x, y):
with backprop.GradientTape() as tape:
z = math_ops.mul(x, array_ops.zeros_like(x))
l = math_ops.add(z, y)
l = math_ops.reduce_sum(l)
gx, gy = tape.gradient(l, [x, y])
x.assign_add(gx)
y.assign_add(gy)
return x + y
# XLA completely optimizes away the variable reads and
# assignments, so skip the test.
if test_util.is_xla_enabled():
self.skipTest('Not relevant for XLA')
with context.eager_mode():
x = resource_variable_ops.ResourceVariable(
np.random.uniform(size=[2, 2]), dtype=dtypes.float32)
y = resource_variable_ops.ResourceVariable(
np.random.uniform(size=[2, 2]), dtype=dtypes.float32)
with context.collect_graphs(optimized=True) as graphs:
f(x, y).numpy()
self.assertLen(graphs, 1)
assign_count = 0
for node in graphs[0].node:
if node.op == 'AssignAddVariableOp':
self.assertEqual(node.input[0], 'y')
assign_count += 1
# Make sure that the only variable update that remains after
# grappler optimization is that of y.
self.assertEqual(assign_count, 1)
self.assertLen(graphs[0].node, 11)
if __name__ == '__main__':
test.main()
tensorflow/python/grappler/cost_analyzer.cc
+301
View File
@@ -0,0 +1,301 @@
/* 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.
==============================================================================*/
#include "tensorflow/python/grappler/cost_analyzer.h"
#include <algorithm>
#include <cmath>
#include <cstdint>
#include <iomanip>
#include <ios>
#include <map>
#include <ostream>
#include <string>
#include <vector>
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "tensorflow/core/framework/cost_graph.pb.h"
#include "tensorflow/core/grappler/costs/op_performance_data.pb.h"
#include "tensorflow/core/grappler/costs/utils.h"
#include "tensorflow/core/grappler/grappler_item.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/protobuf/config.pb.h"
namespace tensorflow {
namespace grappler {
CostAnalyzer::CostAnalyzer(const GrapplerItem& item, Cluster* cluster,
const std::string& suffix)
: item_(&item),
measure_estimator_(cluster, 10, 0),
analytical_estimator_(cluster, /*use_static_shapes=*/false,
/*use_aggressive_shape_inference=*/true),
suffix_(suffix) {}
absl::Status CostAnalyzer::GenerateReport(std::ostream& os,
bool per_node_report, bool verbose) {
GatherCosts();
PreprocessCosts();
AnalyzeCosts();
PrintAnalysis(os, per_node_report, verbose);
return absl::OkStatus();
}
void CostAnalyzer::PredictCosts(CostEstimator* cost_estimator,
CostGraphDef* cost_graph, int64_t* total_time) {
TF_CHECK_OK(cost_estimator->Initialize(*item_));
RunMetadata run_metadata;
Costs costs;
const absl::Status status =
cost_estimator->PredictCosts(item_->graph, &run_metadata, &costs);
if (cost_graph) {
cost_graph->Swap(run_metadata.mutable_cost_graph());
}
*total_time = costs.execution_time.count();
if (!status.ok()) {
LOG(ERROR) << "Could not estimate the cost for item " << item_->id << ": "
<< status.message();
return;
}
}
void CostAnalyzer::GatherCosts() {
CostGraphDef cost_graph_measured;
PredictCosts(&measure_estimator_, &cost_graph_measured,
&total_time_measured_);
VLOG(1) << "Graph size: " << item_->graph.node_size();
VLOG(1) << "cost_graph_measured size: " << cost_graph_measured.node_size();
CostGraphDef cost_graph_analytical;
PredictCosts(&analytical_estimator_, &cost_graph_analytical,
&total_time_analytical_);
VLOG(1) << "cost_graph_analytical size: "
<< cost_graph_analytical.node_size();
CostGraphDef cost_graph_analytical_filtered;
CostGraphDef cost_graph_measured_filtered;
std::map<std::string, const CostGraphDef_Node*> measured_nodes;
for (const auto& node : cost_graph_measured.node()) {
measured_nodes[node.name()] = &node;
}
for (const auto& node : cost_graph_analytical.node()) {
auto it = measured_nodes.find(node.name());
// Filter the nodes that are not the cost nodes returned by
// MeasuringCostEstimator.
if (it == measured_nodes.end()) {
continue;
}
auto added_node_analytical = cost_graph_analytical_filtered.add_node();
auto added_node_measured = cost_graph_measured_filtered.add_node();
*added_node_analytical = node;
*added_node_measured = *(it->second);
}
VLOG(1) << "cost_graph_analytical_filtered size: "
<< cost_graph_analytical_filtered.node_size();
// TODO(yaozhang): add a test to make sure that op_perf_analytical_ and
// op_perf_ cover the same set of nodes.
op_perf_analytical_ = CostGraphToOpPerformanceData(
cost_graph_analytical_filtered, item_->graph);
op_perf_ =
CostGraphToOpPerformanceData(cost_graph_measured_filtered, item_->graph);
}
void CostAnalyzer::PreprocessCosts() {
for (int i = 0; i < op_perf_.op_performance_size(); i++) {
OpPerformance* perf = op_perf_.mutable_op_performance(i);
const OpPerformance& analytical = op_perf_analytical_.op_performance(i);
perf->set_compute_time(analytical.compute_time());
perf->set_memory_time(analytical.memory_time());
double measured_cost = perf->compute_cost();
double analytical_compute_cost = analytical.compute_time();
if (analytical_compute_cost == 0) {
// Negative infinity indicates unavailable data.
perf->set_compute_efficiency(-INFINITY);
} else {
perf->set_compute_efficiency(analytical_compute_cost / measured_cost);
}
double analytical_memory_cost = analytical.memory_time();
if (analytical_memory_cost == 0) {
// Negative infinity indicates unavailable data.
perf->set_memory_efficiency(-INFINITY);
} else {
perf->set_memory_efficiency(analytical_memory_cost / measured_cost);
}
}
}
void CostAnalyzer::SortOpsByTime(std::map<std::string, OpPerfSummary> ops) {
for (const auto& op : ops) {
ops_.push_back(op.second);
}
struct CompareByTime {
bool operator()(const OpPerfSummary& a, const OpPerfSummary& b) const {
return a.time > b.time;
}
};
std::stable_sort(ops_.begin(), ops_.end(), CompareByTime());
}
void CostAnalyzer::AnalyzeCosts() {
std::map<std::string, OpPerfSummary> ops;
for (const auto& op_perf : op_perf_.op_performance()) {
std::string op_name = op_perf.op().op();
ops[op_name].count++;
ops[op_name].time += op_perf.compute_cost();
ops[op_name].compute_time += op_perf.compute_time();
ops[op_name].memory_time += op_perf.memory_time();
ops[op_name].time_upper += op_perf.compute_time() + op_perf.memory_time();
ops[op_name].time_lower +=
std::max(op_perf.compute_time(), op_perf.memory_time());
ops[op_name].name = op_name;
}
SortOpsByTime(ops);
total_time_measured_serialized_ = 0;
total_time_analytical_upper_ = 0;
total_time_analytical_lower_ = 0;
for (const auto& op : ops_) {
total_time_measured_serialized_ += op.time;
total_time_analytical_upper_ += op.time_upper;
total_time_analytical_lower_ += op.time_lower;
}
}
void CostAnalyzer::PrintAnalysis(std::ostream& os, bool per_node_report,
bool verbose) const {
os << std::endl;
os << std::left << std::setw(50)
<< "Total time measured in ns (serialized): " << std::right
<< std::setw(20) << total_time_measured_serialized_ << std::endl;
os << std::left << std::setw(50)
<< "Total time measured in ns (actual): " << std::right << std::setw(20)
<< total_time_measured_ << std::endl;
os << std::left << std::setw(50)
<< "Total time analytical in ns (upper bound): " << std::right
<< std::setw(20) << total_time_analytical_upper_ << std::endl;
os << std::left << std::setw(50)
<< "Total time analytical in ns (lower bound): " << std::right
<< std::setw(20) << total_time_analytical_lower_ << std::endl;
double efficiency_upper = static_cast<double>(total_time_analytical_upper_) /
static_cast<double>(total_time_measured_);
os << std::left << std::setw(50)
<< "Overall efficiency (analytical upper/actual): " << std::right
<< std::setw(20) << efficiency_upper << std::endl;
double efficiency_lower = static_cast<double>(total_time_analytical_lower_) /
static_cast<double>(total_time_measured_);
os << std::left << std::setw(50)
<< "Overall efficiency (analytical lower/actual): " << std::right
<< std::setw(20) << efficiency_lower << std::endl;
os << std::endl;
int width = 35;
int width_narrow = 15;
int width_wide = 20;
os << std::setw(width + 1) << "Op,";
os << std::setw(width_narrow + 1) << "Count,";
os << std::setw(width_wide + 1) << "Measured time (ns),";
os << std::setw(width_narrow + 2) << "Time percent,";
os << std::setw(width_narrow + 2) << "Acc percent,";
os << std::setw(width_wide + 1) << "Analytical upper,";
os << std::setw(width_wide + 1) << "Analytical lower,";
os << std::setw(width_narrow + 2) << "Overall eff";
os << std::setw(width_narrow + 2) << "Compute eff";
os << std::setw(width_narrow + 2) << "Memory eff" << std::endl;
float acc_percent = 0;
for (const auto& op : ops_) {
double percent = static_cast<double>(op.time) /
static_cast<double>(total_time_measured_serialized_);
double eff =
static_cast<double>(op.time_upper) / static_cast<double>(op.time);
double compute_eff =
static_cast<double>(op.compute_time) / static_cast<double>(op.time);
double memory_eff =
static_cast<double>(op.memory_time) / static_cast<double>(op.time);
os << std::setw(width) << op.name << ",";
os << std::setw(width_narrow) << op.count << ",";
os << std::setw(width_wide) << op.time << ",";
os << std::setw(width_narrow) << std::setprecision(2) << percent * 100
<< "%,";
acc_percent += percent;
os << std::setw(width_narrow) << std::setprecision(2) << acc_percent * 100
<< "%,";
os << std::setw(width_wide) << op.time_upper << ",";
os << std::setw(width_wide) << op.time_lower << ",";
os << std::setw(width_narrow) << std::setprecision(2) << eff * 100 << "%,";
os << std::setw(width_narrow) << std::setprecision(2) << compute_eff * 100
<< "%,";
os << std::setw(width_narrow) << std::setprecision(2) << memory_eff * 100
<< "%,";
os << std::endl;
}
os << std::endl;
if (per_node_report) {
if (verbose) {
os << "Below is the full per-node report:" << std::endl;
os << op_perf_.DebugString();
} else {
os << "Below is the per-node report summary:" << std::endl;
int width = 35;
int width_narrow = 15;
int width_wide = 20;
os << std::setw(width + 1) << "Op,";
os << std::setw(width_wide + 1) << "Measured time (ns),";
os << std::setw(width_wide + 1) << "Compute time (ns),";
os << std::setw(width_wide + 1) << "Memory time (ns),";
os << std::setw(width_narrow + 2) << "Compute eff,";
os << std::setw(width_narrow + 2) << "Memory eff,";
os << " Inputs" << std::endl;
for (int i = 0; i < op_perf_.op_performance_size(); i++) {
const auto& perf = op_perf_.op_performance(i);
std::string op_name = perf.op().op();
os << std::setw(width) << op_name << ",";
os << std::setw(width_wide) << perf.compute_cost() << ",";
os << std::setw(width_wide) << perf.compute_time() << ",";
os << std::setw(width_wide) << perf.memory_time() << ",";
os << std::setw(width_narrow) << std::setprecision(2)
<< perf.compute_efficiency() * 100 << "%,";
os << std::setw(width_narrow) << std::setprecision(2)
<< perf.memory_efficiency() * 100 << "%,";
os << " [";
for (int j = 0; j < perf.op().inputs_size(); j++) {
const auto& shape = perf.op().inputs(j).shape();
if (shape.dim_size() > 0) {
os << "(";
std::vector<int> dims;
for (int k = 0; k < shape.dim_size(); k++) {
os << shape.dim(k).size();
if (k < shape.dim_size() - 1) {
os << ", ";
}
}
os << ")";
if (j < perf.op().inputs_size() - 1) {
os << ", ";
}
}
}
os << "]" << std::endl;
}
os << std::endl;
}
}
}
} // end namespace grappler
} // end namespace tensorflow
tensorflow/python/grappler/cost_analyzer.h
+86
View File
@@ -0,0 +1,86 @@
/* 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.
==============================================================================*/
#ifndef TENSORFLOW_PYTHON_GRAPPLER_COST_ANALYZER_H_
#define TENSORFLOW_PYTHON_GRAPPLER_COST_ANALYZER_H_
#include <iostream>
#include "absl/status/status.h"
#include "tensorflow/core/framework/cost_graph.pb.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/tensor_shape.pb.h"
#include "tensorflow/core/grappler/clusters/cluster.h"
#include "tensorflow/core/grappler/costs/analytical_cost_estimator.h"
#include "tensorflow/core/grappler/costs/cost_estimator.h"
#include "tensorflow/core/grappler/costs/measuring_cost_estimator.h"
#include "tensorflow/core/grappler/costs/op_performance_data.pb.h"
namespace tensorflow {
class GraphDef;
class CostGraphDef;
namespace grappler {
struct GrapplerItem;
// Aggregated perf summary for ops of the same type in a graph.
struct OpPerfSummary {
std::string name;
int64_t count;
int64_t time;
int64_t compute_time;
int64_t memory_time;
// Upper and lower bound for estimated time.
int64_t time_upper;
int64_t time_lower;
};
// Generate op-level performance insights on compute/memory
// efficiency, as well as graph-level aggregated performance statistics.
class CostAnalyzer {
public:
explicit CostAnalyzer(const GrapplerItem& item, Cluster* cluster,
const std::string& suffix);
absl::Status GenerateReport(std::ostream& os, bool per_node_report,
bool verbose);
private:
void PredictCosts(CostEstimator* cost_estimator, CostGraphDef* cost_graph,
int64_t* total_time);
void GatherCosts();
void PreprocessCosts();
void AnalyzeCosts();
void SortOpsByTime(std::map<std::string, OpPerfSummary> ops);
void PrintAnalysis(std::ostream& os, bool per_node_report,
bool verbose) const;
const GrapplerItem* item_;
MeasuringCostEstimator measure_estimator_;
AnalyticalCostEstimator analytical_estimator_;
OpPerformanceList op_perf_;
OpPerformanceList op_perf_analytical_;
int64_t total_time_measured_;
int64_t total_time_analytical_;
std::vector<OpPerfSummary> ops_;
int64_t total_time_measured_serialized_;
int64_t total_time_analytical_upper_;
int64_t total_time_analytical_lower_;
std::string suffix_;
};
} // end namespace grappler
} // end namespace tensorflow
#endif // TENSORFLOW_PYTHON_GRAPPLER_COST_ANALYZER_H_
tensorflow/python/grappler/cost_analyzer.py
+81
View File
@@ -0,0 +1,81 @@
# 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.
# =============================================================================
"""Provides a proper python API for the symbols exported through swig."""
from tensorflow.python.grappler import _pywrap_cost_analyzer as tf_wrap
from tensorflow.python.grappler import cluster as gcluster
from tensorflow.python.grappler import item as gitem
def GenerateCostReport(metagraph,
per_node_report=False,
verbose=False,
cluster=None):
"""Analyze the cost of each TensorFlow op and node in the provided metagraph.
Args:
metagraph: A TensorFlow MetaGraphDef.
per_node_report: by default the report contains stats aggregated on a per op
type basis, setting per_node_report to True adds results for each
individual node to the report.
verbose: Prints out the entire operation proto instead of a summary table.
cluster: Analyze the costs using the specified cluster, or the local machine
if no cluster was specified.
Returns:
A string of cost report.
"""
if cluster is None:
cluster = gcluster.Cluster(disable_detailed_stats=False)
return tf_wrap.GenerateCostReport(metagraph.SerializeToString(),
per_node_report, verbose,
cluster.tf_cluster)
def GenerateMemoryReport(metagraph, detailed_report=True, cluster=None):
"""Analyze the peak memory usage for the provided metagraph.
Args:
metagraph: A TensorFlow MetaGraphDef.
detailed_report: print the live tensors in addition to the peak memory
usage.
cluster: Analyze the memory using the specified cluster, or the local
machine if no cluster was specified.
Returns:
A string with the formatted memory usage.
"""
if cluster is None:
cluster = gcluster.Cluster(
disable_detailed_stats=True, disable_timeline=True)
item = gitem.Item(metagraph)
peak_usage = cluster.DeterminePeakMemoryUsage(item)
report = ""
for device, snapshot in peak_usage.items():
peak_usage = snapshot[0]
report += "Peak usage for device " + device + ": " + str(
peak_usage) + " bytes\n"
if detailed_report:
live_tensors = snapshot[1]
for tensor in live_tensors:
op_name = tensor[0]
output_id = tensor[1]
mem_used = tensor[2]
report += " " + str(op_name) + ":" + str(output_id) + " uses " + str(
mem_used) + " bytes\n"
return report
tensorflow/python/grappler/cost_analyzer_test.py
+170
View File
@@ -0,0 +1,170 @@
# 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 the cost analyzer."""
import re
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import meta_graph
from tensorflow.python.framework import ops
from tensorflow.python.framework import test_util
from tensorflow.python.grappler import cost_analyzer
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nn_grad # pylint: disable=unused-import
from tensorflow.python.ops import nn_ops
from tensorflow.python.ops import random_ops
from tensorflow.python.ops import variables
from tensorflow.python.platform import test
from tensorflow.python.training import adam
class CostAnalysisTest(test.TestCase):
@test_util.run_deprecated_v1
def testBasicCost(self):
"""Make sure arguments can be passed correctly."""
a = constant_op.constant(10, name="a")
b = constant_op.constant(20, name="b")
c = math_ops.add_n([a, b], name="c")
d = math_ops.add_n([b, c], name="d")
train_op = ops.get_collection_ref(ops.GraphKeys.TRAIN_OP)
train_op.append(d)
mg = meta_graph.create_meta_graph_def(graph=ops.get_default_graph())
report = cost_analyzer.GenerateCostReport(mg, per_node_report=True)
# Check the report headers
self.assertTrue(b"Total time measured in ns (serialized):" in report)
self.assertTrue(b"Total time measured in ns (actual):" in report)
self.assertTrue(b"Total time analytical in ns (upper bound):" in report)
self.assertTrue(b"Total time analytical in ns (lower bound):" in report)
self.assertTrue(b"Overall efficiency (analytical upper/actual):" in report)
self.assertTrue(b"Overall efficiency (analytical lower/actual):" in report)
self.assertTrue(b"Below is the per-node report summary:" in report)
# Also print the report to make it easier to debug
print("{}".format(report))
@test_util.run_deprecated_v1
def testVerbose(self):
"""Make sure the full report is generated with verbose=True."""
a = constant_op.constant(10, name="a")
b = constant_op.constant(20, name="b")
c = math_ops.add_n([a, b], name="c")
d = math_ops.add_n([b, c], name="d")
train_op = ops.get_collection_ref(ops.GraphKeys.TRAIN_OP)
train_op.append(d)
mg = meta_graph.create_meta_graph_def(graph=ops.get_default_graph())
report = cost_analyzer.GenerateCostReport(
mg, per_node_report=True, verbose=True)
# Check the report headers
self.assertTrue(b"Below is the full per-node report:" in report)
# Also print the report to make it easier to debug
print("{}".format(report))
@test_util.run_deprecated_v1
def testSmallNetworkCost(self):
image = array_ops.placeholder(dtypes.float32, shape=[1, 28, 28, 1])
label = array_ops.placeholder(dtypes.float32, shape=[1, 10])
w = variables.Variable(
random_ops.truncated_normal([5, 5, 1, 32], stddev=0.1))
b = variables.Variable(random_ops.truncated_normal([32], stddev=0.1))
conv = nn_ops.conv2d(image, w, strides=[1, 1, 1, 1], padding="SAME")
h_conv = nn_ops.relu(conv + b)
h_conv_flat = array_ops.reshape(h_conv, [1, -1])
w_fc = variables.Variable(
random_ops.truncated_normal([25088, 10], stddev=0.1))
b_fc = variables.Variable(random_ops.truncated_normal([10], stddev=0.1))
y_conv = nn_ops.softmax(math_ops.matmul(h_conv_flat, w_fc) + b_fc)
cross_entropy = math_ops.reduce_mean(
-math_ops.reduce_sum(label * math_ops.log(y_conv), axis=[1]))
_ = adam.AdamOptimizer(1e-4).minimize(cross_entropy)
mg = meta_graph.create_meta_graph_def(graph=ops.get_default_graph())
report = cost_analyzer.GenerateCostReport(mg)
# Print the report to make it easier to debug
print("{}".format(report))
self.assertTrue(b"MatMul" in report)
self.assertTrue(b"ApplyAdam" in report)
self.assertTrue(b"Conv2DBackpropFilter" in report)
self.assertTrue(b"Softmax" in report)
# When mkl is enabled, Conv2D and MatMul op followed by
# 1-dimension Add in this graph will be fused, but not
# in the mkl disabled case.
expected_matmul_count = 2
op_types = [b"MatMul", b"Conv2DBackpropFilter"]
if not test_util.IsMklEnabled():
self.assertTrue(b"Conv2D" in report)
expected_matmul_count = 3
op_types.append(b"Conv2D")
for op_type in op_types:
matcher = re.compile(
br"\s+" + op_type + br",\s*(\d+),\s*(\d+),\s*([\d\.eE+-]+)%,\s*" +
br"([\d\.eE+-]+)%,\s*(-?\d+),\s*(\d+),", re.MULTILINE)
m = matcher.search(report)
op_count = int(m.group(1))
# upper = int(m.group(5))
lower = int(m.group(6))
if op_type == b"MatMul":
self.assertEqual(expected_matmul_count, op_count)
else:
self.assertEqual(1, op_count)
self.assertTrue(0 <= lower)
# self.assertTrue(0 < upper)
# self.assertTrue(lower <= upper)
@test_util.run_deprecated_v1
def testBasicMemory(self):
"""Make sure arguments can be passed correctly."""
with test_util.device(use_gpu=False):
a = constant_op.constant(10, name="a")
b = constant_op.constant(20, name="b")
c = math_ops.add_n([a, b], name="c")
d = math_ops.add_n([b, c], name="d")
train_op = ops.get_collection_ref(ops.GraphKeys.TRAIN_OP)
train_op.append(d)
mg = meta_graph.create_meta_graph_def(graph=ops.get_default_graph())
report = cost_analyzer.GenerateMemoryReport(mg)
# Print the report to make it easier to debug
print("{}".format(report))
# Check the report
self.assertTrue(
"Peak usage for device /job:localhost/replica:0/task:0/device:CPU:0: "
"16 bytes"
in report)
self.assertTrue(" a:0 uses 4 bytes" in report)
self.assertTrue(" b:0 uses 4 bytes" in report)
self.assertTrue(" c:0 uses 4 bytes" in report)
self.assertTrue(" d:0 uses 4 bytes" in report)
if __name__ == "__main__":
test.main()
tensorflow/python/grappler/cost_analyzer_tool.py
+130
View File
@@ -0,0 +1,130 @@
# 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.
# =============================================================================
"""A tool for cost analysis."""
import argparse
import sys
from absl import app
from google.protobuf import message
from google.protobuf import text_format
from tensorflow.core.framework import graph_pb2
from tensorflow.core.protobuf import config_pb2
from tensorflow.core.protobuf import meta_graph_pb2
from tensorflow.core.protobuf import saved_model_pb2
from tensorflow.python.framework import importer
from tensorflow.python.framework import ops
from tensorflow.python.grappler import cost_analyzer
from tensorflow.python.grappler import tf_optimizer
from tensorflow.python.platform import gfile
from tensorflow.python.training import saver
def get_metagraph():
"""Constructs and returns a MetaGraphDef from the input file."""
with gfile.GFile(FLAGS.input) as input_file:
input_data = input_file.read()
try:
saved_model = saved_model_pb2.SavedModel()
text_format.Merge(input_data, saved_model)
meta_graph = saved_model.meta_graphs[0]
except text_format.ParseError:
try:
saved_model.ParseFromString(input_data)
meta_graph = saved_model.meta_graphs[0]
except message.DecodeError:
try:
meta_graph = meta_graph_pb2.MetaGraphDef()
text_format.Merge(input_data, meta_graph)
except text_format.ParseError:
try:
meta_graph.ParseFromString(input_data)
except message.DecodeError:
try:
graph_def = graph_pb2.GraphDef()
text_format.Merge(input_data, graph_def)
except text_format.ParseError:
try:
graph_def.ParseFromString(input_data)
except message.DecodeError:
raise ValueError(f"Invalid input file: {FLAGS.input}.")
importer.import_graph_def(graph_def, name="")
graph = ops.get_default_graph()
meta_graph = saver.export_meta_graph(
graph_def=graph.as_graph_def(), graph=graph)
if FLAGS.fetch is not None:
fetch_collection = meta_graph_pb2.CollectionDef()
for fetch in FLAGS.fetch.split(","):
fetch_collection.node_list.value.append(fetch)
meta_graph.collection_def["train_op"].CopyFrom(fetch_collection)
return meta_graph
def main(_):
metagraph = get_metagraph()
config = config_pb2.ConfigProto()
if FLAGS.rewriter_config is not None:
text_format.Merge(FLAGS.rewriter_config,
config.graph_options.rewrite_options)
optimized_graph = tf_optimizer.OptimizeGraph(config, metagraph)
metagraph.graph_def.CopyFrom(optimized_graph)
report = cost_analyzer.GenerateCostReport(metagraph, FLAGS.per_node_report,
FLAGS.verbose)
print(report)
if FLAGS.memory_report:
report = cost_analyzer.GenerateMemoryReport(metagraph)
print(report)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--input",
type=str,
default=None,
help="Input file path. Accept SavedModel, MetaGraphDef, and GraphDef in "
"either binary or text format.")
parser.add_argument(
"--fetch",
type=str,
default=None,
help="The names of the fetch node delimited by comma.")
parser.add_argument(
"--rewriter_config",
type=str,
default=None,
help="Configuration for the grappler optimizers, described as a "
"RewriterConfig protocol buffer. Usage example 1: "
"--rewriter_config='optimize_tensor_layout: true "
"disable_model_pruning: true'. Usage example 2: "
"--rewriter_config='optimizers: \"constfold\" optimizers: \"layout\"'")
parser.add_argument(
"--per_node_report",
action="store_true",
help="Generate per-node report. By default the report contains stats "
"aggregated on a per op type basis, per_node_report adds results "
"for each individual node to the report.")
parser.add_argument(
"--memory_report",
action="store_true",
help="Generate memory usage report.")
parser.add_argument(
"--verbose",
action="store_true",
help="Generate verbose reports. By default, succinct reports are used.")
FLAGS, unparsed = parser.parse_known_args()
app.run(main=main, argv=[sys.argv[0]] + unparsed)
tensorflow/python/grappler/cost_analyzer_wrapper.cc
+58
View File
@@ -0,0 +1,58 @@
/* 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 <memory>
#include <sstream>
#include <string>
#include "pybind11/pybind11.h" // from @pybind11
#include "tensorflow/core/grappler/clusters/single_machine.h"
#include "tensorflow/core/grappler/grappler_item.h"
#include "tensorflow/core/grappler/grappler_item_builder.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
#include "tensorflow/python/grappler/cost_analyzer.h"
#include "tensorflow/python/lib/core/pybind11_status.h"
namespace py = pybind11;
PYBIND11_MODULE(_pywrap_cost_analyzer, m) {
m.def("GenerateCostReport",
[](const py::bytes& serialized_metagraph, bool per_node_report,
bool verbose, tensorflow::grappler::Cluster* cluster) -> py::bytes {
tensorflow::MetaGraphDef metagraph;
if (!metagraph.ParseFromString(std::string(serialized_metagraph))) {
return "The MetaGraphDef could not be parsed as a valid protocol "
"buffer";
}
tensorflow::grappler::ItemConfig cfg;
cfg.apply_optimizations = false;
std::unique_ptr<tensorflow::grappler::GrapplerItem> item =
tensorflow::grappler::GrapplerItemFromMetaGraphDef(
"metagraph", metagraph, cfg);
if (item == nullptr) {
return "Error: failed to preprocess metagraph: check your log file "
"for errors";
}
std::string suffix;
tensorflow::grappler::CostAnalyzer analyzer(*item, cluster, suffix);
std::stringstream os;
tensorflow::MaybeRaiseFromStatus(
analyzer.GenerateReport(os, per_node_report, verbose));
return py::bytes(os.str());
});
}
tensorflow/python/grappler/datasets_test.py
+444
View File
@@ -0,0 +1,444 @@
# 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 the datasets shape inference."""
import numpy as np
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.data.ops import iterator_ops
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import meta_graph
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.grappler import item
from tensorflow.python.ops import array_ops
from tensorflow.python.platform import test
class GrapplerTest(test.TestCase):
def testFromTensors(self):
test_cases = [{
'tensor': 0,
'shape': tensor_shape.TensorShape([])
}, {
'tensor': np.array([1, 2, 3]),
'shape': tensor_shape.TensorShape([3])
}, {
'tensor': np.array([[1, 2, 3]]),
'shape': tensor_shape.TensorShape([1, 3])
}]
for test_case in test_cases:
with ops.Graph().as_default() as g:
dataset = dataset_ops.Dataset.from_tensors(test_case['tensor'])
iterator = dataset_ops.make_one_shot_iterator(dataset)
get_next = iterator.get_next()
train_op = ops.get_collection_ref(ops.GraphKeys.TRAIN_OP)
train_op.append(get_next)
mg = meta_graph.create_meta_graph_def(graph=g)
grappler_item = item.Item(mg)
op_properties = grappler_item.GetOpProperties()
self.assertEqual(test_case['shape'],
op_properties['IteratorGetNext'][0].shape)
def testFromTensorSlices(self):
test_cases = [{
'tensor': np.array([1, 2, 3]),
'shape': tensor_shape.TensorShape([])
}, {
'tensor': np.array([[1, 2, 3]]),
'shape': tensor_shape.TensorShape([3])
}, {
'tensor': np.array([[[1, 2, 3]]]),
'shape': tensor_shape.TensorShape([1, 3])
}]
for test_case in test_cases:
with ops.Graph().as_default() as g:
dataset = dataset_ops.Dataset.from_tensor_slices(test_case['tensor'])
iterator = dataset_ops.make_one_shot_iterator(dataset)
get_next = iterator.get_next()
train_op = ops.get_collection_ref(ops.GraphKeys.TRAIN_OP)
train_op.append(get_next)
mg = meta_graph.create_meta_graph_def(graph=g)
grappler_item = item.Item(mg)
op_properties = grappler_item.GetOpProperties()
self.assertEqual(test_case['shape'],
op_properties['IteratorGetNext'][0].shape)
def testFromGenerator(self):
test_cases = [{
'tensor': 0,
'shape': tensor_shape.TensorShape([])
}, {
'tensor': np.array([1, 2, 3]),
'shape': tensor_shape.TensorShape([3])
}, {
'tensor': np.array([[1, 2, 3]]),
'shape': tensor_shape.TensorShape([1, 3])
}]
for test_case in test_cases:
def make_generator(tensor):
def generator():
yield tensor
return generator
with ops.Graph().as_default() as g:
dataset = dataset_ops.Dataset.from_generator(
make_generator(test_case['tensor']),
dtypes.int64,
output_shapes=test_case['shape'])
iterator = dataset_ops.make_one_shot_iterator(dataset)
get_next = iterator.get_next()
train_op = ops.get_collection_ref(ops.GraphKeys.TRAIN_OP)
train_op.append(get_next)
mg = meta_graph.create_meta_graph_def(graph=g)
grappler_item = item.Item(mg)
op_properties = grappler_item.GetOpProperties()
self.assertEqual(test_case['shape'],
op_properties['IteratorGetNext'][0].shape)
def testRange(self):
with ops.Graph().as_default() as g:
dataset = dataset_ops.Dataset.range(42)
iterator = dataset_ops.make_one_shot_iterator(dataset)
get_next = iterator.get_next()
train_op = ops.get_collection_ref(ops.GraphKeys.TRAIN_OP)
train_op.append(get_next)
mg = meta_graph.create_meta_graph_def(graph=g)
grappler_item = item.Item(mg)
op_properties = grappler_item.GetOpProperties()
self.assertEqual(
tensor_shape.TensorShape([]),
op_properties['IteratorGetNext'][0].shape)
def _testTransformation(self, fn):
test_cases = [{
'tensor': 0,
'shape': tensor_shape.TensorShape({})
}, {
'tensor': np.array([1, 2, 3]),
'shape': tensor_shape.TensorShape([3])
}, {
'tensor': np.array([[1, 2, 3]]),
'shape': tensor_shape.TensorShape([1, 3])
}]
for test_case in test_cases:
with ops.Graph().as_default() as g:
dataset = dataset_ops.Dataset.from_tensors(test_case['tensor'])
dataset = fn(dataset, test_case['tensor'], test_case['shape'])
iterator = dataset_ops.make_one_shot_iterator(dataset)
get_next = iterator.get_next()
train_op = ops.get_collection_ref(ops.GraphKeys.TRAIN_OP)
train_op.append(get_next)
mg = meta_graph.create_meta_graph_def(graph=g)
grappler_item = item.Item(mg)
op_properties = grappler_item.GetOpProperties()
self.assertEqual(test_case['shape'],
op_properties['IteratorGetNext'][0].shape)
def testConcatenate(self):
def fn(dataset, tensor, shape):
del shape
return dataset.concatenate(dataset_ops.Dataset.from_tensors(tensor))
self._testTransformation(fn)
def testPrefetch(self):
def fn(dataset, tensor, shape):
del tensor, shape
return dataset.prefetch(42)
self._testTransformation(fn)
def testRepeat(self):
def fn(dataset, tensor, shape):
del tensor, shape
return dataset.repeat(42)
self._testTransformation(fn)
def testShuffle(self):
def fn(dataset, tensor, shape):
del tensor, shape
return dataset.shuffle(42)
self._testTransformation(fn)
def testCache(self):
def fn(dataset, tensor, shape):
del tensor, shape
return dataset.cache()
self._testTransformation(fn)
def testTake(self):
def fn(dataset, tensor, shape):
del tensor, shape
return dataset.take(42)
self._testTransformation(fn)
def testSkip(self):
def fn(dataset, tensor, shape):
del tensor, shape
return dataset.skip(42)
self._testTransformation(fn)
def testShard(self):
def fn(dataset, tensor, shape):
del tensor, shape
return dataset.shard(42, 0)
self._testTransformation(fn)
def testFilter(self):
def fn(dataset, tensor, shape):
del tensor, shape
return dataset.filter(lambda x: True)
self._testTransformation(fn)
def as_tensor_shape(self, proto_with_symbolic_values):
for i in range(len(proto_with_symbolic_values.dim)):
if proto_with_symbolic_values.dim[i].size < -1:
proto_with_symbolic_values.dim[i].size = -1
return tensor_shape.TensorShape(proto_with_symbolic_values)
def testBatch(self):
test_cases = [{
'tensor': 0,
'shape': tensor_shape.TensorShape([None])
}, {
'tensor': np.array([1, 2, 3]),
'shape': tensor_shape.TensorShape([None, 3])
}, {
'tensor': np.array([[1, 2, 3]]),
'shape': tensor_shape.TensorShape([None, 1, 3])
}]
for test_case in test_cases:
with ops.Graph().as_default() as g:
dataset = dataset_ops.Dataset.from_tensors(test_case['tensor'])
dataset = dataset.batch(42)
iterator = dataset_ops.make_one_shot_iterator(dataset)
get_next = iterator.get_next()
train_op = ops.get_collection_ref(ops.GraphKeys.TRAIN_OP)
train_op.append(get_next)
mg = meta_graph.create_meta_graph_def(graph=g)
grappler_item = item.Item(mg)
op_properties = grappler_item.GetOpProperties()
inferred_shape = self.as_tensor_shape(
op_properties['IteratorGetNext'][0].shape)
self.assertTrue(test_case['shape'].dims[0].is_compatible_with(
inferred_shape[0]))
self.assertEqual(test_case['shape'][1:], inferred_shape[1:])
def testPaddedBatch(self):
test_cases = [{
'tensor': 0,
'shape': tensor_shape.TensorShape([None])
}, {
'tensor': np.array([1, 2, 3]),
'shape': tensor_shape.TensorShape([None, 4])
}, {
'tensor': np.array([[1, 2, 3]]),
'shape': tensor_shape.TensorShape([None, 2, 4])
}]
for test_case in test_cases:
with ops.Graph().as_default() as g:
dataset = dataset_ops.Dataset.from_tensors(test_case['tensor'])
dataset = dataset.padded_batch(42, padded_shapes=test_case['shape'][1:])
iterator = dataset_ops.make_one_shot_iterator(dataset)
get_next = iterator.get_next()
train_op = ops.get_collection_ref(ops.GraphKeys.TRAIN_OP)
train_op.append(get_next)
mg = meta_graph.create_meta_graph_def(graph=g)
grappler_item = item.Item(mg)
op_properties = grappler_item.GetOpProperties()
inferred_shape = self.as_tensor_shape(
op_properties['IteratorGetNext'][0].shape)
self.assertTrue(test_case['shape'].dims[0].is_compatible_with(
inferred_shape[0]))
self.assertEqual(test_case['shape'][1:], inferred_shape[1:])
def testFlatMap(self):
test_cases = [{
'tensor': 0,
'shape': tensor_shape.TensorShape([])
}, {
'tensor': np.array([1, 2, 3]),
'shape': tensor_shape.TensorShape([3])
}, {
'tensor': np.array([[1, 2, 3]]),
'shape': tensor_shape.TensorShape([1, 3])
}]
for test_case in test_cases:
with ops.Graph().as_default() as g:
dataset = dataset_ops.Dataset.range(42)
def make_dataset(tensor):
def dataset_fn(n):
return dataset_ops.Dataset.from_tensors(tensor).repeat(n)
return dataset_fn
dataset = dataset.flat_map(make_dataset(test_case['tensor']))
iterator = dataset_ops.make_one_shot_iterator(dataset)
get_next = iterator.get_next()
train_op = ops.get_collection_ref(ops.GraphKeys.TRAIN_OP)
train_op.append(get_next)
mg = meta_graph.create_meta_graph_def(graph=g)
grappler_item = item.Item(mg)
op_properties = grappler_item.GetOpProperties()
self.assertEqual(test_case['shape'],
op_properties['IteratorGetNext'][0].shape)
def testInterleave(self):
test_cases = [{
'tensor': 0,
'shape': tensor_shape.TensorShape([])
}, {
'tensor': np.array([1, 2, 3]),
'shape': tensor_shape.TensorShape([3])
}, {
'tensor': np.array([[1, 2, 3]]),
'shape': tensor_shape.TensorShape([1, 3])
}]
for test_case in test_cases:
with ops.Graph().as_default() as g:
dataset = dataset_ops.Dataset.range(42)
def make_dataset(tensor):
def dataset_fn(n):
return dataset_ops.Dataset.from_tensors(tensor).repeat(n)
return dataset_fn
dataset = dataset.interleave(
make_dataset(test_case['tensor']), cycle_length=42)
iterator = dataset_ops.make_one_shot_iterator(dataset)
get_next = iterator.get_next()
train_op = ops.get_collection_ref(ops.GraphKeys.TRAIN_OP)
train_op.append(get_next)
mg = meta_graph.create_meta_graph_def(graph=g)
grappler_item = item.Item(mg)
op_properties = grappler_item.GetOpProperties()
self.assertEqual(test_case['shape'],
op_properties['IteratorGetNext'][0].shape)
def testMap(self):
test_cases = [{
'tensor': 0,
'shape': tensor_shape.TensorShape([])
}, {
'tensor': np.array([1, 2, 3]),
'shape': tensor_shape.TensorShape([3])
}, {
'tensor': np.array([[1, 2, 3]]),
'shape': tensor_shape.TensorShape([3, 1])
}, {
'tensor': np.array([[[1, 2, 3], [4, 5, 6]]]),
'shape': tensor_shape.TensorShape([3, 2, 1])
}]
for test_case in test_cases:
with ops.Graph().as_default() as g:
dataset = dataset_ops.Dataset.from_tensors(test_case['tensor'])
dataset = dataset.map(array_ops.transpose)
iterator = dataset_ops.make_one_shot_iterator(dataset)
get_next = iterator.get_next()
train_op = ops.get_collection_ref(ops.GraphKeys.TRAIN_OP)
train_op.append(get_next)
mg = meta_graph.create_meta_graph_def(graph=g)
grappler_item = item.Item(mg)
op_properties = grappler_item.GetOpProperties()
self.assertEqual(test_case['shape'],
op_properties['IteratorGetNext'][0].shape)
def testFromStructure(self):
test_cases = [{
'shape': tensor_shape.TensorShape([])
}, {
'shape': tensor_shape.TensorShape([3])
}, {
'shape': tensor_shape.TensorShape([1, 2])
}, {
'shape': tensor_shape.TensorShape([1, 2, 3])
}]
for test_case in test_cases:
with ops.Graph().as_default() as g:
iterator = iterator_ops.Iterator.from_structure(
dtypes.int64, output_shapes=test_case['shape'])
get_next = iterator.get_next()
train_op = ops.get_collection_ref(ops.GraphKeys.TRAIN_OP)
train_op.append(get_next)
mg = meta_graph.create_meta_graph_def(graph=g)
grappler_item = item.Item(mg)
op_properties = grappler_item.GetOpProperties()
self.assertEqual(test_case['shape'],
op_properties['IteratorGetNext'][0].shape)
def testFromStringHandle(self):
test_cases = [{
'shape': tensor_shape.TensorShape([])
}, {
'shape': tensor_shape.TensorShape([3])
}, {
'shape': tensor_shape.TensorShape([1, 2])
}, {
'shape': tensor_shape.TensorShape([1, 2, 3])
}]
for test_case in test_cases:
with ops.Graph().as_default() as g:
iterator = iterator_ops.Iterator.from_structure(dtypes.int64)
handle = iterator.string_handle()
iterator = iterator_ops.Iterator.from_string_handle(
handle, dtypes.int64, output_shapes=test_case['shape'])
get_next = iterator.get_next()
train_op = ops.get_collection_ref(ops.GraphKeys.TRAIN_OP)
train_op.append(get_next)
mg = meta_graph.create_meta_graph_def(graph=g)
grappler_item = item.Item(mg)
op_properties = grappler_item.GetOpProperties()
self.assertEqual(test_case['shape'],
op_properties['IteratorGetNext'][0].shape)
if __name__ == '__main__':
test.main()
tensorflow/python/grappler/graph_analyzer.py
+43
View File
@@ -0,0 +1,43 @@
# 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.
# =============================================================================
"""A tool that finds all subgraphs of a given size in a TF graph.
The subgraph patterns are sorted by occurrence, and only the transitive fanin
part of the graph with regard to the fetch nodes is considered.
"""
import argparse
import sys
from absl import app
from tensorflow.python.grappler import _pywrap_graph_analyzer as tf_wrap
def main(_):
tf_wrap.GraphAnalyzer(FLAGS.input, FLAGS.n)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--input",
type=str,
default=None,
help="Input file path for a TensorFlow MetaGraphDef.")
parser.add_argument(
"--n", type=int, default=None, help="The size of the subgraphs.")
FLAGS, unparsed = parser.parse_known_args()
app.run(main=main, argv=[sys.argv[0]] + unparsed)
tensorflow/python/grappler/graph_analyzer_tool_wrapper.cc
+22
View File
@@ -0,0 +1,22 @@
/* 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 "tensorflow/core/grappler/graph_analyzer/graph_analyzer_tool.h"
PYBIND11_MODULE(_pywrap_graph_analyzer, m) {
m.def("GraphAnalyzer",
&tensorflow::grappler::graph_analyzer::GraphAnalyzerTool);
}
tensorflow/python/grappler/item.py
+90
View File
@@ -0,0 +1,90 @@
# 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.
# ==============================================================================
"""A python interface for Grappler items."""
from tensorflow.core.grappler.costs import op_performance_data_pb2
from tensorflow.core.protobuf import meta_graph_pb2
from tensorflow.python.grappler import _pywrap_tf_item as tf_item
class Item(object):
"""GrapplerItem."""
def __init__(self,
metagraph,
ignore_colocation=True,
ignore_user_placement=False):
"""Creates an Item.
Args:
metagraph: a TensorFlow metagraph.
ignore_colocation: if set, the tool will ignore all the colocation
constraints generated by TensorFlow.
ignore_user_placement: if set, all the placement annotations annotated in
the metagraph will be ignored.
Raises:
ValueError: the metagraph is incomplete or invalid.
"""
self._metagraph = metagraph
self._item_graph = meta_graph_pb2.MetaGraphDef()
self._item_graph.CopyFrom(metagraph)
self._ignore_colocation = ignore_colocation
self._ignore_user_placement = ignore_user_placement
self._tf_item = None
self._BuildTFItem()
def IdentifyImportantOps(self, sort_topologically=False):
return tf_item.TF_IdentifyImportantOps(self.tf_item, sort_topologically)
def GetOpProperties(self):
"""Get Op properties."""
props = tf_item.TF_GetOpProperties(self.tf_item)
properties = {}
for key, values in props.items():
prop = []
for value in values:
# TODO(petebu): Make this conversion to a dictionary be done in the C++
# wrapper for performance.
prop.append(
op_performance_data_pb2.OpInfo.TensorProperties.FromString(value))
properties[key] = prop
return properties
def GetColocationGroups(self):
"""Return a list of hard colocation constraints.
All the nodes in a colocation tuple must be placed on the same device for
the model to work.
Returns:
A list of colocation tuples.
"""
return tf_item.TF_GetColocationGroups(self.tf_item)
@property
def metagraph(self):
return self._metagraph
@property
def tf_item(self):
if self._item_graph != self._metagraph:
self._BuildTFItem()
self._item_graph.CopyFrom(self._metagraph)
return self._tf_item
def _BuildTFItem(self):
self._tf_item = tf_item.TF_NewItem(self._metagraph.SerializeToString(),
self._ignore_colocation,
self._ignore_user_placement)
tensorflow/python/grappler/item_test.py
+144
View File
@@ -0,0 +1,144 @@
# 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 the pybind11 wrapper of Grappler items."""
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import errors_impl
from tensorflow.python.framework import meta_graph
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import test_util
from tensorflow.python.grappler import item
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import gen_array_ops
from tensorflow.python.ops import state_ops
from tensorflow.python.ops import variable_v1
from tensorflow.python.platform import test
class ItemTest(test.TestCase):
"""Unit tests for Grappler Item pybind11 wrapper functionality."""
def _create_sample_metagraph(self, include_train_op=True):
with ops.Graph().as_default() as g:
a = constant_op.constant(10)
b = constant_op.constant(20)
c = a + b
z = control_flow_ops.no_op()
if include_train_op:
train_op = ops.get_collection_ref(ops.GraphKeys.TRAIN_OP)
train_op.append(c)
return meta_graph.create_meta_graph_def(graph=g), z
def test_invalid_item_missing_train_op_raises(self):
"""Verifies that Item raises InvalidArgumentError when train_op is missing."""
mg, _ = self._create_sample_metagraph(include_train_op=False)
with self.assertRaisesRegex(
errors_impl.InvalidArgumentError,
'train_op not specified in the metagraph'):
item.Item(mg)
def test_important_ops_identification(self):
"""Verifies important ops are correctly identified from the metagraph."""
mg, _ = self._create_sample_metagraph(include_train_op=True)
grappler_item = item.Item(mg)
op_list = grappler_item.IdentifyImportantOps()
self.assertCountEqual(['Const', 'Const_1', 'add'], op_list)
def test_op_properties_extraction(self):
"""Verifies op properties are correctly extracted for graph nodes."""
mg, z = self._create_sample_metagraph(include_train_op=True)
grappler_item = item.Item(mg)
op_properties = grappler_item.GetOpProperties()
z_prop = op_properties[z.name]
self.assertEmpty(z_prop)
const_prop = op_properties['Const']
self.assertLen(const_prop, 1)
self.assertEqual(dtypes.int32, const_prop[0].dtype)
self.assertEqual(tensor_shape.TensorShape([]), const_prop[0].shape)
def test_tf_item_initial_properties_equal(self):
"""Verifies initial tf_item properties are consistent and equal."""
mg, _ = self._create_sample_metagraph(include_train_op=True)
grappler_item = item.Item(mg)
initial_tf_item = grappler_item.tf_item
no_change_tf_item = grappler_item.tf_item
self.assertEqual(initial_tf_item, no_change_tf_item)
def test_tf_item_device_modification_updates_wrapper(self):
"""Verifies modifying node placement creates a new underlying tf_item."""
mg, _ = self._create_sample_metagraph(include_train_op=True)
grappler_item = item.Item(mg)
initial_tf_item = grappler_item.tf_item
for node in grappler_item.metagraph.graph_def.node:
node.device = '/cpu:0'
new_tf_item = grappler_item.tf_item
self.assertNotEqual(initial_tf_item, new_tf_item)
def test_tf_item_identical_device_reassignment_unchanged(self):
"""Verifies re-assigning identical placement keeps tf_item unchanged."""
mg, _ = self._create_sample_metagraph(include_train_op=True)
grappler_item = item.Item(mg)
for node in grappler_item.metagraph.graph_def.node:
node.device = '/cpu:0'
new_tf_item = grappler_item.tf_item
for node in grappler_item.metagraph.graph_def.node:
node.device = '/cpu:0'
newest_tf_item = grappler_item.tf_item
self.assertEqual(new_tf_item, newest_tf_item)
@test_util.run_v1_only('b/120545219')
def test_colocation_constraints(self):
"""Verifies colocation constraints are correctly grouped."""
with ops.Graph().as_default() as g:
c = constant_op.constant([10])
v = variable_v1.VariableV1([3], dtype=dtypes.int32)
i = gen_array_ops.ref_identity(v)
a = state_ops.assign(i, c)
train_op = ops.get_collection_ref(ops.GraphKeys.TRAIN_OP)
train_op.append(a)
mg = meta_graph.create_meta_graph_def(graph=g)
grappler_item = item.Item(mg)
groups = grappler_item.GetColocationGroups()
self.assertLen(groups, 1)
self.assertCountEqual(
groups[0], ['Assign', 'RefIdentity', 'Variable', 'Variable/Assign'])
@test_util.run_v1_only('b/120545219')
def test_colocation_constraints_missing_input(self):
"""Verifies standalone nodes with missing inputs are correctly grouped."""
with ops.Graph().as_default() as g:
c = constant_op.constant([10])
v = variable_v1.VariableV1([3], dtype=dtypes.int32)
i = gen_array_ops.ref_identity(v)
a = state_ops.assign(i, c)
train_op = ops.get_collection_ref(ops.GraphKeys.TRAIN_OP)
train_op.append(a)
mg = meta_graph.create_meta_graph_def(graph=g)
for node in mg.graph_def.node:
if node.op == 'Assign':
del node.input[:]
grappler_item = item.Item(mg)
groups = grappler_item.GetColocationGroups()
self.assertLen(groups, 1)
self.assertCountEqual(
groups[0], ['RefIdentity', 'Variable'])
if __name__ == '__main__':
test.main()
tensorflow/python/grappler/item_wrapper.cc
+263
View File
@@ -0,0 +1,263 @@
/* 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 <algorithm>
#include <memory>
#include <stdexcept>
#include <string>
#include <unordered_map>
#include <vector>
#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "pybind11/pybind11.h" // from @pybind11
#include "pybind11/stl.h" // from @pybind11
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/grappler/costs/graph_properties.h"
#include "tensorflow/core/grappler/costs/op_performance_data.pb.h"
#include "tensorflow/core/grappler/grappler_item.h"
#include "tensorflow/core/grappler/grappler_item_builder.h"
#include "tensorflow/core/grappler/utils.h"
#include "tensorflow/core/grappler/utils/topological_sort.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
#include "tensorflow/python/lib/core/pybind11_status.h"
namespace py = pybind11;
// Manages disjoint sets of colocation groups using a union-find data structure
// with path compression and union by rank.
class ColocationGroups {
public:
// Ensures a node is tracked in the disjoint-set structure.
void RegisterNode(absl::string_view node_name) { Find(node_name); }
// Unions the colocation sets containing nodes x and y.
void Group(absl::string_view x, absl::string_view y) {
Rep* x_root = Find(x);
Rep* y_root = Find(y);
// x and y are already in the same set
if (x_root == y_root) {
return;
}
// x and y are not in same set, so we merge them
// Use the occasion to strengthen what we know about the handle by merging
// the information about the 2 subsets.
if (x_root->rank < y_root->rank) {
x_root->parent = y_root;
} else if (x_root->rank > y_root->rank) {
y_root->parent = x_root;
} else {
// Arbitrarily make one root the new parent
y_root->parent = x_root;
++x_root->rank;
}
}
// Extracts all disjoint colocation groups as a list of node name lists.
std::vector<std::vector<std::string>> ExtractGroups() {
std::vector<std::vector<std::string>> groups;
groups.reserve(nodes_.size());
absl::flat_hash_map<const Rep*, int> group_ids;
for (const auto& [node_name, rep_ptr] : nodes_) {
Rep* r = Find(rep_ptr.get());
auto [it, inserted] = group_ids.try_emplace(r, groups.size());
int id = it->second;
if (inserted) {
// If inserted, this is a new group. The value stored (groups.size())
// is the index where the new group should be added.
groups.emplace_back();
}
groups[id].push_back(node_name);
}
for (auto& g : groups) {
std::sort(g.begin(), g.end());
}
std::sort(groups.begin(), groups.end());
return groups;
}
private:
struct Rep {
// Parent in the tree used to encode the set.
Rep* parent;
// Rank in the tree, used to figure out how to compress the path to the root
// of the tree.
int rank;
};
Rep* Find(absl::string_view n) {
// Try to emplace a new Rep. If the key already exists, try_emplace does
// nothing and returns an iterator to the existing element. Otherwise,
// it inserts a new unique_ptr<Rep> and returns an iterator to it.
auto [it, inserted] = nodes_.try_emplace(n, std::make_unique<Rep>());
if (inserted) {
// First time processing this handle, initialize the new entry.
Rep* raw_node = it->second.get();
raw_node->parent = raw_node;
raw_node->rank = 0;
return raw_node;
}
return Find(it->second.get());
}
Rep* Find(Rep* node) {
Rep* root = node->parent;
while (root != root->parent) {
root = root->parent;
}
while (node->parent != root) {
Rep* next = node->parent;
node->parent = root;
node = next;
}
return root;
}
absl::flat_hash_map<std::string, std::unique_ptr<Rep>> nodes_;
};
PYBIND11_MAKE_OPAQUE(tensorflow::grappler::GrapplerItem);
PYBIND11_MODULE(_pywrap_tf_item, m) {
py::class_<tensorflow::grappler::GrapplerItem> grappler_item(m,
"GrapplerItem");
m.def("TF_NewItem",
[](const py::bytes& serialized_metagraph, bool ignore_colocation,
bool ignore_user_placement)
-> std::unique_ptr<tensorflow::grappler::GrapplerItem> {
tensorflow::MetaGraphDef metagraph;
py::buffer_info info = py::buffer(serialized_metagraph).request();
if (!metagraph.ParseFromArray(info.ptr, info.size)) {
throw std::invalid_argument(
"The MetaGraphDef could not be parsed as a valid protocol "
"buffer");
}
if (metagraph.collection_def().count("train_op") == 0) {
tsl::MaybeRaiseRegisteredFromStatus(absl::InvalidArgumentError(
"train_op not specified in the metagraph"));
}
tensorflow::grappler::ItemConfig cfg;
cfg.ignore_user_placement = ignore_user_placement;
cfg.ignore_colocation = ignore_colocation;
std::unique_ptr<tensorflow::grappler::GrapplerItem> item =
tensorflow::grappler::GrapplerItemFromMetaGraphDef(
"item", metagraph, cfg);
if (item == nullptr) {
tsl::MaybeRaiseRegisteredFromStatus(
absl::InvalidArgumentError("Invalid metagraph"));
}
return item;
});
m.def("TF_IdentifyImportantOps",
[](tensorflow::grappler::GrapplerItem* item,
bool sort_topologically) -> std::vector<std::string> {
std::vector<const tensorflow::NodeDef*> main_ops =
item->MainOpsFanin();
std::vector<const tensorflow::NodeDef*> enqueue_ops =
item->EnqueueOpsFanin();
absl::flat_hash_set<std::string> op_names;
for (const tensorflow::NodeDef* op : main_ops) {
op_names.insert(op->name());
}
for (const tensorflow::NodeDef* op : enqueue_ops) {
op_names.insert(op->name());
}
std::vector<std::string> ops;
if (sort_topologically) {
tensorflow::GraphDef subgraph;
for (const tensorflow::NodeDef& node : item->graph.node()) {
if (op_names.find(node.name()) != op_names.end()) {
*subgraph.add_node() = node;
}
}
tsl::MaybeRaiseRegisteredFromStatus(
tensorflow::grappler::TopologicalSort(&subgraph));
for (const tensorflow::NodeDef& node : subgraph.node()) {
ops.push_back(node.name());
}
} else {
ops.reserve(op_names.size());
for (const auto& op_name : op_names) {
ops.push_back(op_name);
}
std::sort(ops.begin(), ops.end());
}
return ops;
});
m.def("TF_GetOpProperties",
[](tensorflow::grappler::GrapplerItem* item)
-> std::unordered_map<std::string, std::vector<py::bytes>> {
tensorflow::grappler::GraphProperties properties(*item);
tsl::MaybeRaiseRegisteredFromStatus(
properties.InferStatically(false));
std::unordered_map<std::string, std::vector<py::bytes>> props;
for (const tensorflow::NodeDef& node : item->graph.node()) {
const std::string& node_name = node.name();
const std::vector<tensorflow::OpInfo::TensorProperties>&
output_props = properties.GetOutputProperties(node_name);
std::vector<py::bytes> prop;
prop.reserve(output_props.size());
for (const tensorflow::OpInfo::TensorProperties& output_prop :
output_props) {
prop.push_back(output_prop.SerializeAsString());
}
props[node_name] = std::move(prop);
}
return props;
});
m.def("TF_GetColocationGroups",
[](tensorflow::grappler::GrapplerItem* item)
-> std::vector<std::vector<std::string>> {
ColocationGroups groupings;
tensorflow::OpRegistry* registry = tensorflow::OpRegistry::Global();
for (const tensorflow::NodeDef& node : item->graph.node()) {
const tensorflow::OpDef* op_def;
if (!registry->LookUpOpDef(node.op(), &op_def).ok()) {
continue;
}
tensorflow::NameRangeMap inputs;
if (!tensorflow::NameRangesForNode(node, *op_def, &inputs, nullptr)
.ok()) {
continue;
}
for (const tensorflow::OpDef::ArgDef& arg : op_def->input_arg()) {
if (!arg.is_ref()) {
continue;
}
const auto& range = inputs[arg.name()];
for (int i = range.first;
i < range.second && i < node.input_size(); ++i) {
groupings.Group(node.name(),
tensorflow::grappler::NodeName(node.input(i)));
}
}
}
return groupings.ExtractGroups();
});
}
tensorflow/python/grappler/layout_optimizer_test.py
+2121
View File
File diff suppressed because it is too large Load Diff
tensorflow/python/grappler/memory_optimizer_test.py
+326
View File
@@ -0,0 +1,326 @@
# 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 the swig wrapper tf_optimizer."""
from tensorflow.core.framework import attr_value_pb2
from tensorflow.core.protobuf import config_pb2
from tensorflow.core.protobuf import rewriter_config_pb2
from tensorflow.python.client import session
from tensorflow.python.framework import meta_graph
from tensorflow.python.framework import ops
from tensorflow.python.framework import random_seed
from tensorflow.python.framework import test_util
from tensorflow.python.grappler import tf_optimizer
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nn
from tensorflow.python.ops import variable_scope
from tensorflow.python.ops import variable_v1
from tensorflow.python.ops import variables
from tensorflow.python.platform import test
from tensorflow.python.training import training as train
class MemoryOptimizerSwapTest(test.TestCase):
"""Tests the Grappler memory optimizer."""
@test_util.run_deprecated_v1
def testNoSwapping(self):
"""Make sure the graph is preserved when there is nothing to swap."""
a = variable_v1.VariableV1(10, name='a')
b = variable_v1.VariableV1(20, name='b')
c = math_ops.add_n([a, b], name='c')
d = math_ops.add_n([b, c], name='d')
train_op = ops.get_collection_ref(ops.GraphKeys.TRAIN_OP)
train_op.append(d)
mg = meta_graph.create_meta_graph_def(graph=ops.get_default_graph())
graph_size = len(mg.graph_def.node)
nodes = [node.name for node in mg.graph_def.node]
config = config_pb2.ConfigProto()
config.graph_options.rewrite_options.CopyFrom(
rewriter_config_pb2.RewriterConfig(
disable_model_pruning=True,
constant_folding=rewriter_config_pb2.RewriterConfig.OFF,
dependency_optimization=rewriter_config_pb2.RewriterConfig.OFF,
memory_optimization=rewriter_config_pb2.RewriterConfig.MANUAL))
graph = tf_optimizer.OptimizeGraph(config, mg)
self.assertEqual(len(graph.node), graph_size)
self.assertItemsEqual([node.name for node in graph.node], nodes)
@test_util.run_v1_only('b/120545219')
def testSimpleSwap(self):
"""Check that the swap annotations are followed."""
with ops.device('/gpu:0'):
a = variable_v1.VariableV1(10, name='a')
b = variable_v1.VariableV1(20, name='b')
c = math_ops.add_n([a, b], name='c')
d = math_ops.add_n([b, c], name='d')
train_op = ops.get_collection_ref(ops.GraphKeys.TRAIN_OP)
train_op.append(d)
d.op._set_attr('_swap_to_host', attr_value_pb2.AttrValue(i=0))
mg = meta_graph.create_meta_graph_def(graph=ops.get_default_graph())
graph_size = len(mg.graph_def.node)
config = config_pb2.ConfigProto()
config.graph_options.rewrite_options.CopyFrom(
rewriter_config_pb2.RewriterConfig(
disable_model_pruning=True,
meta_optimizer_iterations=rewriter_config_pb2.RewriterConfig.ONE,
constant_folding=rewriter_config_pb2.RewriterConfig.OFF,
memory_optimization=rewriter_config_pb2.RewriterConfig.MANUAL,
min_graph_nodes=-1))
graph = tf_optimizer.OptimizeGraph(config, mg)
self.assertEqual(len(graph.node), graph_size + 2)
self.assertTrue(
set(node.name for node in graph.node) > set(
['a', 'b', 'c', 'd', 'swap_in_d_0', 'swap_out_d_0']))
for node in graph.node:
if node.name == 'swap_in_d_0':
self.assertEqual('swap_out_d_0', node.input[0])
self.assertEqual('^b/read', node.input[1])
elif node.name == 'swap_out_d_0':
self.assertEqual('b/read', node.input[0])
elif node.name == 'd':
self.assertEqual('swap_in_d_0', node.input[0])
self.assertEqual('c', node.input[1])
class MemoryOptimizerRecomputeTest(test.TestCase):
"""Tests the Python interface to recomputation rewrites.
See core/grappler/optimizers/memory_optimizer_test.cc for functional tests.
"""
def _GetMetaGraph(self, batch_size=14, image_dim=12, optimizer_scope_name=''):
"""A simple layered graph with conv, an intermediate op, and a ReLU."""
graph = ops.Graph()
with graph.as_default():
random_seed.set_random_seed(1)
current_activation = variable_scope.get_variable(
name='start', shape=[batch_size, image_dim, image_dim, 5])
conv_filter = variable_scope.get_variable(
name='filter', shape=[5, 5, 5, 5])
for layer_number in range(10):
with variable_scope.variable_scope('layer_{}'.format(layer_number)):
after_conv = nn.conv2d(current_activation, conv_filter, [1, 1, 1, 1],
'SAME')
current_activation = 2. * after_conv
current_activation = nn.relu(current_activation)
loss = math_ops.reduce_mean(current_activation)
with ops.name_scope(optimizer_scope_name):
optimizer = train.AdamOptimizer(0.001)
train_op = optimizer.minimize(loss)
init_op = variables.global_variables_initializer()
metagraph = train.export_meta_graph()
return (metagraph, init_op.name, train_op.name, loss.name)
def testRewritingDefaultGradientNames(self):
"""Tests that rewriting occurs with default gradient names."""
(original_metagraph, _, _, _) = self._GetMetaGraph()
config = config_pb2.ConfigProto()
config.graph_options.rewrite_options.CopyFrom(
rewriter_config_pb2.RewriterConfig(
disable_model_pruning=True,
constant_folding=rewriter_config_pb2.RewriterConfig.OFF,
dependency_optimization=rewriter_config_pb2.RewriterConfig.OFF,
layout_optimizer=rewriter_config_pb2.RewriterConfig.OFF,
arithmetic_optimization=rewriter_config_pb2.RewriterConfig.OFF,
min_graph_nodes=-1,
memory_optimization=(
rewriter_config_pb2.RewriterConfig.RECOMPUTATION_HEURISTICS)))
rewritten_graph_def = tf_optimizer.OptimizeGraph(config, original_metagraph)
self.assertGreater(
len(rewritten_graph_def.node),
len(original_metagraph.graph_def.node))
self.assertEqual(
0,
len([node for node in original_metagraph.graph_def.node
if 'Recomputed/' in node.name]))
self.assertEqual(
20, # Two per layer
len([node for node in rewritten_graph_def.node
if 'Recomputed/' in node.name]))
def testRewritingNameScopedGradientNames(self):
"""Tests that rewriting occurs with non-standard gradient names."""
(original_metagraph, _, _, _) = self._GetMetaGraph(
optimizer_scope_name='optimizer')
config = config_pb2.ConfigProto()
config.graph_options.rewrite_options.CopyFrom(
rewriter_config_pb2.RewriterConfig(
disable_model_pruning=True,
constant_folding=rewriter_config_pb2.RewriterConfig.OFF,
dependency_optimization=rewriter_config_pb2.RewriterConfig.OFF,
layout_optimizer=rewriter_config_pb2.RewriterConfig.OFF,
arithmetic_optimization=rewriter_config_pb2.RewriterConfig.OFF,
min_graph_nodes=-1,
memory_optimization=rewriter_config_pb2.RewriterConfig
.RECOMPUTATION_HEURISTICS,
# Checks that name scope "gradients/" also match sub-scope.
memory_optimizer_target_node_name_scope='gradients/'))
rewritten_graph_def = tf_optimizer.OptimizeGraph(config, original_metagraph)
self.assertGreater(
len(rewritten_graph_def.node),
len(original_metagraph.graph_def.node))
self.assertEqual(
0,
len([node for node in original_metagraph.graph_def.node
if 'Recomputed/' in node.name]))
self.assertEqual(
20, # Two per layer
len([node for node in rewritten_graph_def.node
if 'Recomputed/' in node.name]))
def testRewritingNameScopedGradientNamesScope(self):
"""Tests that rewriting occurs with non-standard gradient names."""
(original_metagraph, _, _,
_) = self._GetMetaGraph(optimizer_scope_name='foo/bar')
config = config_pb2.ConfigProto()
config.graph_options.rewrite_options.CopyFrom(
rewriter_config_pb2.RewriterConfig(
disable_model_pruning=True,
constant_folding=rewriter_config_pb2.RewriterConfig.OFF,
dependency_optimization=rewriter_config_pb2.RewriterConfig.OFF,
layout_optimizer=rewriter_config_pb2.RewriterConfig.OFF,
arithmetic_optimization=rewriter_config_pb2.RewriterConfig.OFF,
memory_optimization=rewriter_config_pb2.RewriterConfig
.RECOMPUTATION_HEURISTICS,
# This should not match anything.
memory_optimizer_target_node_name_scope='r/gradients/'))
rewritten_graph_def = tf_optimizer.OptimizeGraph(config, original_metagraph)
self.assertEqual(
len(rewritten_graph_def.node), len(original_metagraph.graph_def.node))
self.assertEqual(0,
len([
node for node in original_metagraph.graph_def.node
if 'Recomputed/' in node.name
]))
self.assertEqual(0,
len([
node for node in rewritten_graph_def.node
if 'Recomputed/' in node.name
]))
def _GetMemoryOptimizerSessionConfig(self):
rewrite_options = rewriter_config_pb2.RewriterConfig(
disable_model_pruning=True,
memory_optimization=rewriter_config_pb2.RewriterConfig.HEURISTICS)
graph_options = config_pb2.GraphOptions(rewrite_options=rewrite_options)
return config_pb2.ConfigProto(graph_options=graph_options)
def _RunMetaGraphWithConfig(
self, config, metagraph, init_op_name, train_op_name, loss_op_name):
graph = ops.Graph()
with graph.as_default():
train.import_meta_graph(metagraph)
init_op = graph.get_operation_by_name(init_op_name)
train_op = graph.get_operation_by_name(train_op_name)
loss_op = graph.get_tensor_by_name(loss_op_name)
with session.Session(config=config, graph=graph) as sess:
self.evaluate(init_op)
self.evaluate(train_op)
self.evaluate(train_op)
return self.evaluate(loss_op)
def testRecomputationRewritingNoErrors(self):
"""Tests that graph output is not significantly different with rewriting."""
(original_metagraph, init_op_name, train_op_name, loss_op_name
) = self._GetMetaGraph()
original_loss = self._RunMetaGraphWithConfig(
config=config_pb2.ConfigProto(),
metagraph=original_metagraph,
init_op_name=init_op_name,
train_op_name=train_op_name,
loss_op_name=loss_op_name)
memory_optimized_loss = self._RunMetaGraphWithConfig(
config=self._GetMemoryOptimizerSessionConfig(),
metagraph=original_metagraph,
init_op_name=init_op_name,
train_op_name=train_op_name,
loss_op_name=loss_op_name)
self.assertAllClose(original_loss, memory_optimized_loss, rtol=1e-2)
def _annotated_graph(self):
graph = ops.Graph()
with graph.as_default():
random_seed.set_random_seed(2)
current_activation = variable_scope.get_variable(
name='start', shape=[1, 2, 2, 5])
conv_filter = variable_scope.get_variable(
name='filter', shape=[5, 5, 5, 5])
for layer_number in range(3):
with variable_scope.variable_scope('layer_{}'.format(layer_number)):
after_conv = nn.conv2d(current_activation, conv_filter, [1, 1, 1, 1],
'SAME')
current_activation = 2. * after_conv
current_activation.op._set_attr(
'_recompute_hint',
# The value of the attribute does not matter; just that the key
# exists in the op's attributes.
attr_value_pb2.AttrValue(i=1))
current_activation += 5.
current_activation.op._set_attr(
'_recompute_hint', attr_value_pb2.AttrValue(i=0))
current_activation = nn.relu(current_activation)
current_activation.op._set_attr(
'_recompute_hint', attr_value_pb2.AttrValue(i=1))
loss = math_ops.reduce_mean(current_activation)
optimizer = train.AdamOptimizer(0.001)
train_op = optimizer.minimize(loss)
init_op = variables.global_variables_initializer()
return graph, init_op, train_op
def testHintNoMetaGraph(self):
# Closer to expected usage, but does not check that a re-write actually
# happens; see testHintDoesRewrite.
graph, init_op, train_op = self._annotated_graph()
with graph.as_default():
manual_memory_config = rewriter_config_pb2.RewriterConfig(
memory_optimization=rewriter_config_pb2.RewriterConfig.MANUAL)
graph_options = config_pb2.GraphOptions(
rewrite_options=manual_memory_config)
session_config = config_pb2.ConfigProto(graph_options=graph_options)
with session.Session(config=session_config) as sess:
self.evaluate(init_op)
self.evaluate(train_op)
@test_util.run_v1_only('b/120545219')
def testHintDoesRewrite(self):
graph = self._annotated_graph()[0]
with graph.as_default():
metagraph = train.export_meta_graph()
self.assertEqual(
0,
len([node for node in metagraph.graph_def.node
if 'Recomputed/' in node.name]))
config = config_pb2.ConfigProto()
config.graph_options.rewrite_options.CopyFrom(
rewriter_config_pb2.RewriterConfig(
min_graph_nodes=-1,
memory_optimization=rewriter_config_pb2.RewriterConfig.MANUAL))
rewritten_graph_def = tf_optimizer.OptimizeGraph(config, metagraph)
self.assertEqual(
9,
len([
node for node in rewritten_graph_def.node
if 'Recomputed/' in node.name
]))
if __name__ == '__main__':
test.main()
tensorflow/python/grappler/model_analyzer.cc
+110
View File
@@ -0,0 +1,110 @@
/* 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.
==============================================================================*/
#include "tensorflow/python/grappler/model_analyzer.h"
#include <ostream>
#include <vector>
#include "absl/status/status.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/tensor_shape.pb.h"
#include "tensorflow/core/grappler/costs/graph_properties.h"
#include "tensorflow/core/grappler/costs/op_performance_data.pb.h"
#include "tensorflow/core/grappler/grappler_item.h"
namespace tensorflow {
namespace grappler {
ModelAnalyzer::ModelAnalyzer(const GrapplerItem& item) : item_(item) {}
absl::Status ModelAnalyzer::GenerateReport(bool debug, bool assume_valid_feeds,
std::ostream& os) {
GraphProperties properties(item_);
TF_RETURN_IF_ERROR(properties.InferStatically(assume_valid_feeds));
for (const auto& node : item_.MainOpsFanin()) {
PrintNodeInfo(node, properties, debug, os);
}
for (const auto& node : item_.EnqueueOpsFanin()) {
PrintNodeInfo(node, properties, debug, os);
}
return absl::OkStatus();
}
void ModelAnalyzer::PrintNodeInfo(const NodeDef* node,
const GraphProperties& properties, bool debug,
std::ostream& os) const {
os << node->name() << " [" << node->op() << "]" << std::endl;
if (properties.HasOutputProperties(node->name())) {
const std::vector<OpInfo::TensorProperties>& props =
properties.GetOutputProperties(node->name());
for (int i = 0, props_size = props.size(); i < props_size; ++i) {
const OpInfo::TensorProperties& prop = props[i];
os << "\t"
<< "output " << i << " (" << DataTypeString(prop.dtype())
<< ") has shape ";
if (prop.shape().unknown_rank()) {
os << "?";
} else {
os << "[";
for (int i = 0; i < prop.shape().dim_size(); ++i) {
if (i > 0) {
os << ", ";
}
if (prop.shape().dim(i).size() >= 0) {
// Print the actual dimension.
os << prop.shape().dim(i).size();
} else if (prop.shape().dim(i).size() == -1) {
// We don't know anything about the dimension.
os << "?";
} else {
// Symbolic dimension.
os << "x" << -prop.shape().dim(i).size();
}
}
os << "]";
}
os << std::endl;
}
}
if (debug) {
const OpRegistrationData* op_reg_data;
absl::Status status =
OpRegistry::Global()->LookUp(node->op(), &op_reg_data);
if (!status.ok()) {
os << "\tCouldn't find op registration for " << node->op() << std::endl;
} else if (!op_reg_data->shape_inference_fn) {
os << "\tCouldn't find shape function for op " << node->op() << std::endl;
} else if (properties.HasInputProperties(node->name())) {
const std::vector<OpInfo::TensorProperties>& props =
properties.GetInputProperties(node->name());
for (int i = 0, props_size = props.size(); i < props_size; ++i) {
const OpInfo::TensorProperties& prop = props[i];
if (prop.has_value()) {
os << "\t"
<< "input " << i << " (" << DataTypeString(prop.dtype())
<< ") has known value" << std::endl;
}
}
}
}
}
} // end namespace grappler
} // end namespace tensorflow
tensorflow/python/grappler/model_analyzer.h
+49
View File
@@ -0,0 +1,49 @@
/* 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.
==============================================================================*/
#ifndef TENSORFLOW_PYTHON_GRAPPLER_MODEL_ANALYZER_H_
#define TENSORFLOW_PYTHON_GRAPPLER_MODEL_ANALYZER_H_
#include <iostream>
#include "absl/status/status.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/lib/core/status.h"
namespace tensorflow {
namespace grappler {
struct GrapplerItem;
class GraphProperties;
// Generate a report detailing how much information is known statically for most
// operations in the model, including output data types and output shapes.
class ModelAnalyzer {
public:
explicit ModelAnalyzer(const GrapplerItem& item);
absl::Status GenerateReport(bool debug, bool assume_valid_feeds,
std::ostream& os);
private:
void PrintNodeInfo(const NodeDef* node, const GraphProperties& properties,
bool debug, std::ostream& os) const;
const GrapplerItem& item_;
};
} // end namespace grappler
} // end namespace tensorflow
#endif // TENSORFLOW_PYTHON_GRAPPLER_MODEL_ANALYZER_H_
tensorflow/python/grappler/model_analyzer.py
+32
View File
@@ -0,0 +1,32 @@
# 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.
# =============================================================================
"""Provides a proper python API for the symbols exported through swig."""
from tensorflow.python.grappler import _pywrap_model_analyzer as tf_wrap
def GenerateModelReport(metagraph, assume_valid_feeds=True, debug=False):
"""Report what's known statically about each node in the provided metagraph.
Args:
metagraph: A TensorFlow MetaGraphDef.
assume_valid_feeds: If True, assume that the shape of the fed nodes is valid
debug: Add some information useful for debugging.
Returns:
A string containing the report.
"""
return tf_wrap.GenerateModelReport(
metagraph.SerializeToString(), assume_valid_feeds, debug)
tensorflow/python/grappler/model_analyzer_test.py
+71
View File
@@ -0,0 +1,71 @@
# 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 the cost analyzer."""
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import meta_graph
from tensorflow.python.framework import ops
from tensorflow.python.framework import test_util
from tensorflow.python.grappler import model_analyzer
from tensorflow.python.ops import math_ops
from tensorflow.python.platform import test
class PyWrapOptimizeGraphTest(test.TestCase):
@test_util.run_deprecated_v1
def testBasic(self):
"""Make sure arguments can be passed correctly."""
a = constant_op.constant([10, 11], name="a")
b = constant_op.constant([10], name="b")
c = math_ops.add(a, b, name="c")
d = math_ops.add_n([a, c], name="d")
train_op = ops.get_collection_ref(ops.GraphKeys.TRAIN_OP)
train_op.append(d)
mg = meta_graph.create_meta_graph_def(graph=ops.get_default_graph())
report = model_analyzer.GenerateModelReport(mg)
# Check the report headers
self.assertIn(b"a [Const]", report)
self.assertIn(b"a [Const]", report)
self.assertIn(b"c [AddV2]", report)
self.assertIn(b"d [AddN]", report)
# Also print the report to make it easier to debug
print("{}".format(report))
@test_util.run_deprecated_v1
def testDebugMode(self):
"""Make sure arguments can be passed correctly."""
a = constant_op.constant([10, 11], name="a")
b = constant_op.constant([10], name="b")
c = math_ops.add(a, b, name="c")
train_op = ops.get_collection_ref(ops.GraphKeys.TRAIN_OP)
train_op.append(c)
mg = meta_graph.create_meta_graph_def(graph=ops.get_default_graph())
report = model_analyzer.GenerateModelReport(mg, debug=True)
# Check the report headers
self.assertIn(b"input 0 (int32) has known value", report)
self.assertIn(b"input 1 (int32) has known value", report)
# Also print the report to make it easier to debug
print("{}".format(report))
if __name__ == "__main__":
test.main()
tensorflow/python/grappler/model_analyzer_wrapper.cc
+55
View File
@@ -0,0 +1,55 @@
/* 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 <memory>
#include <sstream>
#include <string>
#include "pybind11/pybind11.h" // from @pybind11
#include "tensorflow/core/grappler/grappler_item_builder.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
#include "tensorflow/python/grappler/model_analyzer.h"
#include "tensorflow/python/lib/core/pybind11_status.h"
namespace py = pybind11;
PYBIND11_MODULE(_pywrap_model_analyzer, m) {
m.def("GenerateModelReport",
[](const py::bytes& serialized_metagraph, bool assume_valid_feeds,
bool debug) -> py::bytes {
tensorflow::MetaGraphDef metagraph;
if (!metagraph.ParseFromString(std::string(serialized_metagraph))) {
return "The MetaGraphDef could not be parsed as a valid protocol "
"buffer";
}
tensorflow::grappler::ItemConfig cfg;
cfg.apply_optimizations = false;
std::unique_ptr<tensorflow::grappler::GrapplerItem> item =
tensorflow::grappler::GrapplerItemFromMetaGraphDef(
"metagraph", metagraph, cfg);
if (item == nullptr) {
return "Error: failed to preprocess metagraph: check your log file "
"for errors";
}
tensorflow::grappler::ModelAnalyzer analyzer(*item);
std::ostringstream os;
tensorflow::MaybeRaiseFromStatus(
analyzer.GenerateReport(debug, assume_valid_feeds, os));
return py::bytes(os.str());
});
}
tensorflow/python/grappler/remapper_test.py
+342
View File
@@ -0,0 +1,342 @@
# 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 Grappler Remapper."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os
from absl.testing import parameterized
from tensorflow.core.protobuf import config_pb2
from tensorflow.core.protobuf import rewriter_config_pb2
from tensorflow.python.client import session
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import meta_graph
from tensorflow.python.framework import ops
from tensorflow.python.framework import test_util
from tensorflow.python.grappler import tf_optimizer
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import init_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nn
from tensorflow.python.ops import nn_ops
from tensorflow.python.ops import random_ops
from tensorflow.python.ops import variables
from tensorflow.python.platform import sysconfig as sysconfig_lib
from tensorflow.python.platform import test
from tensorflow.python.util import _pywrap_utils
def _input(shape):
"""Generates an input of a given shape."""
return variables.Variable(random_ops.truncated_normal(shape, seed=0))
def _weight(shape):
"""Generates a weight of a given shape."""
# Note that the lambda is needed to allow construction inside loops.
return variables.Variable(lambda: init_ops.glorot_uniform_initializer(seed=0)
(shape))
def _bias(shape):
"""Generates a bias of a given shape."""
return constant_op.constant(0.1, shape=shape)
def _get_config(remapping_on=False):
"""Returns a CongfigProto with remapper optimizer on/off."""
rewrite_config = rewriter_config_pb2.RewriterConfig(
remapping=rewriter_config_pb2.RewriterConfig
.ON if remapping_on else rewriter_config_pb2.RewriterConfig.OFF)
rewrite_config.min_graph_nodes = -1
graph_options = config_pb2.GraphOptions(rewrite_options=rewrite_config)
config = config_pb2.ConfigProto(graph_options=graph_options)
return config
class RemapperTest(test.TestCase, parameterized.TestCase):
"""Tests the Grappler remapper optimizer."""
def setUp(self):
super(RemapperTest, self).setUp()
# GeluApproximate fusion on GPU requires cublasLt.
os.environ['TF_USE_CUBLASLT'] = '1'
os.environ['TF_CUDNN_USE_RUNTIME_FUSION'] = '1'
def maybe_skip_test(self, mode):
if mode == 'cuda':
# It seems the windows os cannot correctly query the cuda_version.
# TODO(kaixih@nvidia): Remove this when it works.
if os.name == 'nt':
self.skipTest("This test doesn't support Windows")
# The cublaslt matmul with gelu epilog is only supported since cuda 11.4.
if not test.is_gpu_available(cuda_only=True):
self.skipTest('This test requires GPU.')
cuda_version_str = sysconfig_lib.get_build_info().get(
'cuda_version', '0.0')
cuda_version = tuple([int(x) for x in cuda_version_str.split('.')])
if cuda_version < (11, 4):
self.skipTest('This test requires CUDA >= 11.4.')
if mode == 'mkl' and not test_util.IsMklEnabled():
self.skipTest('MKL is not enabled.')
def _VerifyNoFusion(self, model_fn):
ops.add_to_collection('train_op', model_fn)
mg = meta_graph.create_meta_graph_def(graph=model_fn.graph)
# Compute reference
config = _get_config(remapping_on=False)
gdef_ref = tf_optimizer.OptimizeGraph(config, mg)
# Compute with remapping ON
config = _get_config(remapping_on=True)
gdef = tf_optimizer.OptimizeGraph(config, mg)
self.assertEqual(len(gdef_ref.node), len(gdef.node))
self.assertAllEqual([n.op for n in gdef_ref.node],
[n.op for n in gdef.node])
def _VerifyValues(self, model_fn, use_low_precision, fused_op, epilog_ops):
run_options = config_pb2.RunOptions(output_partition_graphs=True)
metadata = config_pb2.RunMetadata()
# Compute reference value.
config = _get_config(remapping_on=False)
with session.Session(config=config) as sess:
sess.run(variables.global_variables_initializer())
output_ref = sess.run(
model_fn, options=run_options, run_metadata=metadata)
# Compute output with fusion.
config = _get_config(remapping_on=True)
with session.Session(config=config) as sess:
sess.run(variables.global_variables_initializer())
output_val = sess.run(
model_fn, options=run_options, run_metadata=metadata)
graph = metadata.partition_graphs[0]
# Graph should contain fused op.
found_fused_op = False
for node in graph.node:
if node.op in fused_op:
fused_ops = node.attr['fused_ops'].list.s
ops_matched = len(fused_ops) >= 1 and len(fused_ops) == len(epilog_ops)
for op_a, op_b in zip(fused_ops, epilog_ops):
if op_a != op_b:
ops_matched = False
break
found_fused_op = ops_matched
break
self.assertTrue(found_fused_op)
# Computed output value should be close to reference value.
tol = 1e-2 if use_low_precision else 1e-5
self.assertAllClose(output_ref, output_val, atol=tol, rtol=tol)
return graph
@parameterized.parameters(['cuda', 'mkl'])
@test_util.run_deprecated_v1
@test_util.disable_xla('This test does not pass with XLA')
@test_util.run_without_tensor_float_32('Avoid TF32 convs on A100+ GPUs')
def test_matmul_biasadd_activation_fusion(self, mode):
"""Test MatMul+BiasAdd+Gelu fusion."""
self.maybe_skip_test(mode)
def gelu_approximate(x):
return nn.gelu(x, approximate=True)
def gelu_exact(x):
return nn.gelu(x, approximate=False)
# Erfc-based implementation of GeluExact from:
# https://github.com/tensorflow/tensorflow/pull/76174
def gelu_exact_erfc(x):
return (
0.5
* x
* math_ops.erfc(-x * math_ops.cast(0.7071067811865476, x.dtype))
)
device = '/device:GPU:0' if mode == 'cuda' else '/device:CPU:0'
config = []
use_fp16 = True
if (
test_util.IsMklEnabled()
and not _pywrap_utils.IsDataTypeSupportedByOneDNNOnThisCPU(
dtypes.float16
)
):
use_fp16 = False
if mode == 'mkl':
config.append((dtypes.float32, gelu_exact, b'GeluExact'))
config.append((dtypes.float32, gelu_exact_erfc, b'GeluExact'))
config.append((dtypes.float32, gelu_approximate, b'GeluApproximate'))
if _pywrap_utils.IsDataTypeSupportedByOneDNNOnThisCPU(dtypes.bfloat16):
config.append((dtypes.bfloat16, gelu_approximate, b'GeluApproximate'))
config.append((dtypes.bfloat16, gelu_exact, b'GeluExact'))
elif mode == 'cuda':
config.append((dtypes.float32, gelu_approximate, b'GeluApproximate'))
if use_fp16:
config.append((dtypes.float16, gelu_approximate, b'GeluApproximate'))
# Gelu exact fusion is supported by cuDNN frontend APIs and performant
# with fp16 and on Ampere GPUs and later.
if (test_util.is_gpu_available(
cuda_only=True, min_cuda_compute_capability=(8, 0))):
config.append((dtypes.float16, gelu_exact, b'GeluExact'))
config.append((dtypes.float16, math_ops.tanh, b'Tanh'))
config.append((dtypes.float16, math_ops.sigmoid, b'Sigmoid'))
m, n, k = (2, 4, 6) # Matrix dimensions
fused_op = ['_MklNativeFusedMatMul', '_MklFusedMatMul', '_FusedMatMul']
for precision, act_fn, act_name in config:
for transpose in (False, True):
# Create MatMul + BiasAdd + Activation graph
ops.reset_default_graph()
with ops.device(device):
x = _input([k, m] if transpose else [m, k])
w = _weight([n, k] if transpose else [k, n])
b = _bias([n])
x = math_ops.cast(x, precision)
w = math_ops.cast(w, precision)
b = math_ops.cast(b, precision)
y = math_ops.matmul(
x, w, transpose_a=transpose, transpose_b=transpose)
z = nn.bias_add(y, b)
out = act_fn(z)
if transpose and (device == '/device:CPU:0') and \
act_name in (b'GeluApproximate', b'GeluExact'):
if precision == dtypes.bfloat16:
# No fusion should happen on CPU.
self._VerifyNoFusion(out)
continue
else:
# Gelu should not get fused, only BiasAdd.
epilog_ops = [b'BiasAdd']
else:
epilog_ops = [b'BiasAdd', act_name]
graph = self._VerifyValues(out, precision != dtypes.float32, fused_op,
epilog_ops)
@test_util.run_deprecated_v1
@test_util.disable_xla('This test does not pass with XLA')
@test_util.run_without_tensor_float_32('Avoid TF32 convs on A100+ GPUs')
def test_conv2d_biasadd_act_fusion(self):
"""Test Conv2D+BiasAdd+Relu fusion."""
if not test_util.is_gpu_available():
self.skipTest('No GPU available')
if test.is_built_with_rocm():
self.skipTest('ROCm does not support conv biasadd fusion')
N, H, W, C = (5, 3, 3, 8) # pylint: disable=invalid-name
# The runtime fusion requires the output dims to be 32-bit aligned.
self.assertEqual(C % 2, 0)
act_fns = [nn.relu]
act_names = [b'Relu']
if test_util.is_gpu_available(
cuda_only=True, min_cuda_compute_capability=(8, 0)):
act_fns += [nn.elu, nn.relu6, nn.leaky_relu]
act_names += [b'Elu', b'Relu6', b'LeakyRelu']
for precision in ('float16', 'float32'):
for act_fn, act_name in zip(act_fns, act_names):
use_fp16 = precision == 'float16'
if (
test_util.IsMklEnabled()
and use_fp16
and not _pywrap_utils.IsDataTypeSupportedByOneDNNOnThisCPU(
dtypes.float16
)
):
continue
# The runtime fusion (when the activation is not relu) only supports
# fp16 at this moment.
if not use_fp16 and act_name != b'Relu':
continue
ops.reset_default_graph()
x_shape = [N, C, H, W]
x_format, b_format = ('NCHW', 'NC..')
if use_fp16:
x_shape = [N, H, W, C]
x_format, b_format = ('NHWC', 'N..C')
x = _input(x_shape)
w = _weight([2, 2, C, C])
b = _bias([C])
if use_fp16:
x = math_ops.cast(x, dtypes.float16)
w = math_ops.cast(w, dtypes.float16)
b = math_ops.cast(b, dtypes.float16)
y = nn_ops.conv2d(
x, w, strides=(1, 1), padding='SAME', data_format=x_format)
z = nn.bias_add(y, b, data_format=b_format)
out = act_fn(z)
out = array_ops.identity(out)
epilog_ops = [b'BiasAdd', act_name]
fused_op = ['_FusedConv2D']
graph = self._VerifyValues(out, use_fp16, fused_op, epilog_ops)
@test_util.run_deprecated_v1
@test_util.disable_xla('This test does not pass with XLA')
@test_util.run_without_tensor_float_32('Avoid TF32 convs on A100+ GPUs')
def test_two_conv2d_fusions(self):
"""Test two Conv2D patterns and only the second is fusable."""
if not test_util.is_gpu_available(
cuda_only=True, min_cuda_compute_capability=(8, 0)):
self.skipTest('No GPU with compute compatibility >= 8.0 available')
N, H, W, C = (5, 3, 3, 8) # pylint: disable=invalid-name
ops.reset_default_graph()
x_shape = [N, C, H, W]
x_format, b_format = ('NCHW', 'NC..')
x = _input(x_shape)
w = _weight([2, 2, C, C])
b = _bias([C])
y = nn_ops.conv2d(
x, w, strides=(1, 1), padding='SAME', data_format=x_format)
y = nn.bias_add(y, b, data_format=b_format)
y = nn.leaky_relu(y)
y = nn_ops.conv2d(
y, w, strides=(1, 1), padding='SAME', data_format=x_format)
y = nn.bias_add(y, b, data_format=b_format)
y = nn.relu(y)
out = array_ops.identity(y)
# The first Conv-BiasAdd-LeakyRelu is not fusable because cuDNN requires
# fp16 for this pattern. The second Conv-BiasAdd-Relu is fusable.
epilog_ops = [b'BiasAdd', b'Relu']
fused_op = ['_FusedConv2D']
self._VerifyValues(out, False, fused_op, epilog_ops)
if __name__ == '__main__':
test.main()
tensorflow/python/grappler/tf_optimizer.py
+91
View File
@@ -0,0 +1,91 @@
# 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.
# =============================================================================
"""Provides a proper python API for the symbols exported through swig."""
import threading
from tensorflow.core.framework import graph_pb2
from tensorflow.core.protobuf import config_pb2
from tensorflow.python.grappler import _pywrap_tf_optimizer as tf_opt
from tensorflow.python.grappler import cluster as gcluster
_OPTIMIZE_GRAPH_CLUSTER_LOCK = threading.Lock()
is_oss = True # Updated by copybara.
def OptimizeGraph(config_proto,
metagraph,
verbose=True,
graph_id=b'graph_to_optimize',
cluster=None,
strip_default_attributes=False):
"""Optimize the provided metagraph.
For best results, the signature_def field in `metagraph` should be populated
with information about input (feed) and output (fetch) tensors.
Args:
config_proto: a ConfigProto protobuf.
metagraph: a MetagraphDef protobuf.
verbose: whether to log optimization results.
graph_id: a string identifying this graph.
cluster: a grappler cluster object representing hardware resources
available to run this graph.
strip_default_attributes: whether graph node attributes having default
values should be removed after all the optimization passes. This
option is useful if the resulting graph will be executed by an older
process that might not know some of the recently added attributes.
"""
if not isinstance(config_proto, config_pb2.ConfigProto):
raise TypeError('Argument `config_proto` should be a tf.ConfigProto, '
f'received type: {type(config_proto).__name__}')
if is_oss:
optimize_method = tf_opt.TF_OptimizeGraphSerialized
metagraph = metagraph.SerializeToString()
else:
optimize_method = tf_opt.TF_OptimizeGraph
if cluster is not None:
out_graph = optimize_method(
cluster.tf_cluster,
config_proto.SerializeToString(),
metagraph,
verbose,
graph_id,
strip_default_attributes,
)
else:
# Currently Grappler assumes no more than 1 sessions alive globally.
# See comments on SingleMachine::Provision(), hence we use the following
# lock to prevent concurrent access to the following code.
with _OPTIMIZE_GRAPH_CLUSTER_LOCK:
cluster = gcluster.Cluster()
try:
out_graph = optimize_method(
cluster.tf_cluster,
config_proto.SerializeToString(),
metagraph,
verbose,
graph_id,
strip_default_attributes,
)
finally:
# Force the cleanup instead of waiting on python GC to cleanup the
# temporary cluster we've created. Otherwise subsequent calls might
# not have a clean slate because GC may not have run yet.
cluster.Shutdown()
if is_oss:
out_graph = graph_pb2.GraphDef.FromString(out_graph)
return out_graph
tensorflow/python/grappler/tf_optimizer_test.py
+136
View File
@@ -0,0 +1,136 @@
# 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 the swig wrapper tf_optimizer."""
from tensorflow.core.framework import attr_value_pb2
from tensorflow.core.protobuf import config_pb2
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import meta_graph
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import test_util
from tensorflow.python.grappler import item as gitem
from tensorflow.python.grappler import tf_optimizer
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import variable_v1
from tensorflow.python.ops import while_loop
from tensorflow.python.platform import test
class PyWrapOptimizeGraphTest(test.TestCase):
@test_util.run_deprecated_v1
def testBasic(self):
"""Make sure arguments can be passed correctly."""
a = constant_op.constant(10, name='a')
b = constant_op.constant(20, name='b')
c = math_ops.add_n([a, b], name='c')
d = math_ops.add_n([b, c], name='d')
train_op = ops.get_collection_ref(ops.GraphKeys.TRAIN_OP)
# Being a train_op will make 'd' to be added as a fetch node.
train_op.append(d)
mg = meta_graph.create_meta_graph_def(graph=ops.get_default_graph())
config = config_pb2.ConfigProto()
rewriter_config = config.graph_options.rewrite_options
rewriter_config.optimizers.append('constfold')
rewriter_config.min_graph_nodes = -1
graph = tf_optimizer.OptimizeGraph(config, mg)
self.assertEqual(len(graph.node), 1)
self.assertItemsEqual([node.name for node in graph.node], ['d'])
@test_util.run_v1_only('b/120545219')
def testKeepNodes(self):
g = ops.Graph()
with g.as_default():
a1 = variable_v1.VariableV1(
1.0) # Must be preserved since it's in the collection 'variables'.
a2 = constant_op.constant(0, shape=[50, 50], name='keep')
ops.add_to_collection('a2', a2) # Explicitly add to collection.
with g._attr_scope(
{'_grappler_do_not_remove': attr_value_pb2.AttrValue(b=True)}):
a3 = constant_op.constant(0, name='keep2')
b = constant_op.constant(1, shape=[100, 10])
c = constant_op.constant(0, shape=[10, 30])
d = math_ops.matmul(b, c)
ops.add_to_collection('train_op', d) # d is the fetch node.
# Optimize the graph.
mg = meta_graph.create_meta_graph_def(graph=g)
config = config_pb2.ConfigProto()
rewriter_config = config.graph_options.rewrite_options
rewriter_config.min_graph_nodes = -1
optimized_graph = tf_optimizer.OptimizeGraph(config, mg)
# Check that the nodes referenced in various collections have been preserved
optimized_graph_nodes = [node.name for node in optimized_graph.node]
expected_nodes = [
d.op.name, a1.op.name, a2.op.name, a3.op.name, 'Variable/initial_value',
'Variable/Assign'
]
self.assertEqual(len(optimized_graph_nodes), len(expected_nodes))
self.assertAllInSet(optimized_graph_nodes, expected_nodes)
@test_util.run_v1_only('b/120545219')
def testLoops(self):
g = ops.Graph()
with g.as_default():
def _Cond(_, counter):
return counter < end
def _Body(buf, counter):
buf = array_ops.concat([buf, [counter]], 0)
counter += 1
return [buf, counter]
start = array_ops.placeholder(shape=[], dtype=dtypes.int32)
end = array_ops.placeholder(shape=[], dtype=dtypes.int32)
init_buf = array_ops.zeros(shape=[0], dtype=dtypes.int32)
loop_vars = [init_buf, start]
shape_inv = [
tensor_shape.TensorShape([None]),
tensor_shape.TensorShape([])
]
buf, _ = while_loop.while_loop(_Cond, _Body, loop_vars, shape_inv)
f = -array_ops.ones_like(buf, optimize=False) # pylint: disable=invalid-unary-operand-type
buf_shape = array_ops.shape(buf)
f_shape = array_ops.shape(f)
ops.add_to_collection('train_op', buf_shape)
ops.add_to_collection('train_op', f_shape)
# Optimize the graph.
mg = meta_graph.create_meta_graph_def(graph=g)
config = config_pb2.ConfigProto()
rewriter_config = config.graph_options.rewrite_options
rewriter_config.min_graph_nodes = -1
optimized_graph = tf_optimizer.OptimizeGraph(config, mg)
mg.graph_def.CopyFrom(optimized_graph)
# Check that the nodes referenced in various collections have been preserved
item = gitem.Item(mg)
props = item.GetOpProperties()
buf_prop = props[buf.op.name]
f_prop = props[f.op.name]
self.assertEqual(buf_prop, f_prop)
if __name__ == '__main__':
test.main()
tensorflow/python/grappler/tf_optimizer_wrapper.cc
+143
View File
@@ -0,0 +1,143 @@
/* 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 <memory>
#include <stdexcept>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>
#include "pybind11/pybind11.h" // from @pybind11
#include "pybind11_protobuf/native_proto_caster.h" // from @pybind11_protobuf
#include "tensorflow/core/common_runtime/device.h"
#include "tensorflow/core/common_runtime/device_factory.h"
#include "tensorflow/core/framework/device_attributes.pb.h"
#include "tensorflow/core/framework/device_base.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/graph_def_util.h"
#include "tensorflow/core/grappler/clusters/cluster.h"
#include "tensorflow/core/grappler/clusters/utils.h"
#include "tensorflow/core/grappler/grappler_item.h"
#include "tensorflow/core/grappler/grappler_item_builder.h"
#include "tensorflow/core/grappler/optimizers/meta_optimizer.h"
#include "tensorflow/core/protobuf/config.pb.h"
#include "tensorflow/core/protobuf/device_properties.pb.h"
#include "tensorflow/core/protobuf/meta_graph.pb.h"
#include "tensorflow/core/public/session_options.h"
#include "tensorflow/python/lib/core/pybind11_status.h"
namespace py = pybind11;
void DetectDevices(
std::unordered_map<std::string, tensorflow::DeviceProperties>* device_map) {
tensorflow::SessionOptions options;
std::vector<std::unique_ptr<tensorflow::Device>> devices;
if (!tensorflow::DeviceFactory::AddDevices(options, "", &devices).ok()) {
return;
}
for (const std::unique_ptr<tensorflow::Device>& device : devices) {
tensorflow::DeviceProperties& prop = (*device_map)[device->name()];
prop = tensorflow::grappler::GetDeviceInfo(device->parsed_name());
// Overwrite the memory limit since users might have requested to use only a
// fraction of the available device memory.
const tensorflow::DeviceAttributes& attr = device->attributes();
prop.set_memory_size(attr.memory_limit());
}
}
namespace {
tensorflow::GraphDef TF_OptimizeGraph(
tensorflow::grappler::Cluster* cluster,
const std::string& serialized_config_proto,
const tensorflow::MetaGraphDef& metagraph, bool verbose,
const std::string& graph_id, bool strip_default_attributes) {
tensorflow::GraphDef out_graph;
{
py::gil_scoped_release gil_release;
tensorflow::ConfigProto config_proto;
if (!config_proto.ParseFromString(serialized_config_proto)) {
throw std::invalid_argument(
"The ConfigProto could not be parsed as a valid protocol "
"buffer");
}
tensorflow::grappler::ItemConfig item_config;
// This disables graph optimizations in the older graph optimizer,
// which tend to overlap / be redundant with those in Grappler.
item_config.apply_optimizations = false;
item_config.ignore_user_placement = false;
std::unique_ptr<tensorflow::grappler::GrapplerItem> grappler_item =
tensorflow::grappler::GrapplerItemFromMetaGraphDef(graph_id, metagraph,
item_config);
if (!grappler_item) {
throw std::invalid_argument(
"Failed to import metagraph, check error log for more info.");
}
tensorflow::DeviceBase* cpu_device = nullptr;
tensorflow::grappler::MetaOptimizer optimizer(cpu_device, config_proto);
tsl::MaybeRaiseRegisteredFromStatusWithGIL(
optimizer.Optimize(cluster, *grappler_item, &out_graph));
if (strip_default_attributes) {
tensorflow::StripDefaultAttributes(*tensorflow::OpRegistry::Global(),
out_graph.mutable_node());
}
if (verbose) {
optimizer.PrintResult();
}
}
return out_graph;
}
} // namespace
// Here two bindings of `TF_OptimizeGraph` are introduced, one serializes the
// `MetaGraphDef` to string when passing it to C++, another uses pybind's
// protobuf interface, and enables its native proto caster feature
// (https://github.com/pybind/pybind11_protobuf#c-native-vs-python-native-types)
// to prevent serialization. This is necessary for grappler to support models
// larger than 2GiB.
// At the moment, the open source python API defaults to the serialized
// implementation.
PYBIND11_MODULE(_pywrap_tf_optimizer, m) {
pybind11_protobuf::ImportNativeProtoCasters();
m.def("TF_OptimizeGraphSerialized",
[](tensorflow::grappler::Cluster* cluster,
const std::string& serialized_config_proto,
const std::string& serialized_metagraph, bool verbose,
const std::string& graph_id,
bool strip_default_attributes) -> py::bytes {
std::string out_graph_bytes;
{
tensorflow::MetaGraphDef metagraph;
if (!metagraph.ParseFromString(serialized_metagraph)) {
throw std::invalid_argument(
"The MetaGraphDef could not be parsed as a valid protocol "
"buffer");
}
const tensorflow::GraphDef out_graph =
TF_OptimizeGraph(cluster, serialized_config_proto, metagraph,
verbose, graph_id, strip_default_attributes);
out_graph_bytes = out_graph.SerializeAsString();
}
return py::bytes(std::move(out_graph_bytes));
});
m.def("TF_OptimizeGraph", &TF_OptimizeGraph);
}