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/tools/graph_transforms/BUILD
+353
View File
@@ -0,0 +1,353 @@
# Description:
# Utilities that perform useful transformations on graphs
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
load(
"//tensorflow:tensorflow.bzl",
"if_not_windows",
"tf_cc_binary",
"tf_cc_test",
"tf_copts",
)
load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable", "tf_py_strict_test")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = ["//visibility:public"],
licenses = ["notice"],
)
cc_library(
name = "transform_utils",
srcs = [
"transform_utils.cc",
],
hdrs = [
"transform_utils.h",
],
compatible_with = get_compatible_with_portable(),
copts = tf_copts(),
visibility = ["//visibility:public"],
deps = [
"//tensorflow/core:framework",
"//tensorflow/core:framework_internal",
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
"//tensorflow/core:protos_all_cc",
],
)
tf_cc_test(
name = "transform_utils_test",
size = "small",
srcs = ["transform_utils_test.cc"],
deps = [
":transform_utils",
"//tensorflow/cc:cc_ops",
"//tensorflow/core:core_cpu_base",
"//tensorflow/core:framework",
"//tensorflow/core:framework_internal",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:tensorflow",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core:testlib",
],
)
cc_library(
name = "file_utils",
srcs = [
"file_utils.cc",
],
hdrs = [
"file_utils.h",
],
copts = tf_copts(),
visibility = ["//visibility:public"],
deps = [
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
],
)
tf_cc_test(
name = "file_utils_test",
size = "small",
srcs = ["file_utils_test.cc"],
deps = [
":file_utils",
"//tensorflow/cc:cc_ops",
"//tensorflow/core:core_cpu",
"//tensorflow/core:framework_internal",
"//tensorflow/core:lib",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core:testlib",
],
)
cc_library(
name = "transforms_lib",
srcs = [
"add_default_attributes.cc",
"backports.cc",
"flatten_atrous.cc",
"fold_batch_norms.cc",
"fold_constants_lib.cc",
"fold_old_batch_norms.cc",
"freeze_requantization_ranges.cc",
"fuse_convolutions.cc",
"inline_partitionedcall.cc",
"insert_logging.cc",
"obfuscate_names.cc",
"quantize_nodes.cc",
"quantize_weights.cc",
"remove_attribute.cc",
"remove_control_dependencies.cc",
"remove_device.cc",
"remove_nodes.cc",
"rename_attribute.cc",
"rename_node.cc",
"rename_op.cc",
"round_weights.cc",
"set_device.cc",
"sort_by_execution_order.cc",
"sparsify_gather.cc",
"strip_unused_nodes.cc",
],
hdrs = [
"fold_constants_lib.h",
],
copts = tf_copts(),
visibility = ["//visibility:public"],
deps = [
":transform_utils",
"//tensorflow/c:checkpoint_reader",
"//tensorflow/core:core_cpu",
"//tensorflow/core:core_cpu_internal",
"//tensorflow/core:framework",
"//tensorflow/core:framework_internal",
"//tensorflow/core:graph",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:tensorflow",
"//tensorflow/core/kernels:quantization_utils",
"//tensorflow/core/util/tensor_bundle",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/strings",
] + if_not_windows([
"//tensorflow/core:sparse_ops_op_lib",
"//tensorflow/core:parsing_ops_op_lib",
"//tensorflow/core:sendrecv_ops_op_lib",
"//tensorflow/core:io_ops_op_lib",
"//tensorflow/core:logging_ops_op_lib",
"//tensorflow/core:lookup_ops_op_lib",
"//tensorflow/core:data_flow_ops_op_lib",
"//tensorflow/core:no_op_op_lib",
"//tensorflow/core:state_ops_op_lib",
"//tensorflow/core:user_ops_op_lib",
"//tensorflow/core:training_ops_op_lib",
"//tensorflow/core:string_ops_op_lib",
"//tensorflow/core:random_ops_op_lib",
"//tensorflow/core:rnn_ops_op_lib",
"//tensorflow/core:nn_ops_op_lib",
"//tensorflow/core:math_ops_op_lib",
"//tensorflow/core:manip_ops_op_lib",
"//tensorflow/core:list_ops_op_lib",
"//tensorflow/core:functional_ops_op_lib",
"//tensorflow/core:control_flow_ops_op_lib",
"//tensorflow/core:candidate_sampling_ops_op_lib",
"//tensorflow/core:array_ops_op_lib",
]),
alwayslink = 1,
)
tf_cc_test(
name = "transforms_test",
size = "small",
srcs = [
"add_default_attributes_test.cc",
"backports_test.cc",
"flatten_atrous_test.cc",
"fold_batch_norms_test.cc",
"fold_constants_test.cc",
"fold_old_batch_norms_test.cc",
"freeze_requantization_ranges_test.cc",
"fuse_convolutions_test.cc",
"inline_partitionedcall_test.cc",
"insert_logging_test.cc",
"obfuscate_names_test.cc",
"quantize_nodes_test.cc",
"quantize_weights_test.cc",
"remove_attribute_test.cc",
"remove_device_test.cc",
"remove_nodes_test.cc",
"rename_attribute_test.cc",
"rename_node_test.cc",
"rename_op_test.cc",
"round_weights_test.cc",
"set_device_test.cc",
"sort_by_execution_order_test.cc",
"sparsify_gather_test.cc",
"strip_unused_nodes_test.cc",
],
deps = [
":transform_utils",
":transforms_lib",
"//tensorflow/cc:cc_ops",
"//tensorflow/cc:sendrecv_ops",
"//tensorflow/core:bitwise_ops_op_lib",
"//tensorflow/core:core_cpu",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core:testlib",
"//tensorflow/core/kernels:quantization_utils",
"//tensorflow/core/kernels:quantized_ops",
"//tensorflow/core/util/tensor_bundle",
],
)
filegroup(
name = "transform_graph_hdrs",
srcs = [
"transform_graph.h",
"transform_utils.h",
],
visibility = ["//tensorflow/python/util:__pkg__"],
)
cc_library(
name = "transform_graph_lib",
srcs = ["transform_graph.cc"],
hdrs = ["transform_graph.h"],
copts = tf_copts(),
visibility = ["//visibility:public"],
deps = [
":file_utils",
":transform_utils",
":transforms_lib",
"//tensorflow/core:framework_internal",
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
"//tensorflow/core:protos_all_cc",
],
alwayslink = 1,
)
# This library includes a main function, to make it easy to create other
# versions of the tool linked against different operator libs.
cc_library(
name = "transform_graph_main_lib",
srcs = ["transform_graph_main.cc"],
copts = tf_copts(),
visibility = ["//visibility:public"],
deps = [
":transform_graph_lib",
":transforms_lib",
"//tensorflow/core:framework_internal",
"//tensorflow/core:lib",
],
)
tf_cc_binary(
name = "transform_graph",
copts = tf_copts(),
linkstatic = 1,
visibility = ["//visibility:public"],
deps = [
":transform_graph_main_lib",
],
)
tf_cc_test(
name = "transform_graph_test",
size = "medium",
srcs = ["transform_graph_test.cc"],
deps = [
":transform_graph_lib",
":transform_utils",
"//tensorflow/cc:cc_ops",
"//tensorflow/cc:sendrecv_ops",
"//tensorflow/core:core_cpu",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core:testlib",
],
)
# This library includes a main function, to make it easy to create other
# versions of the tool linked against different operator libs.
cc_library(
name = "summarize_graph_main_lib",
srcs = ["summarize_graph_main.cc"],
copts = tf_copts(),
visibility = ["//visibility:public"],
deps = [
":file_utils",
":transform_utils",
"//tensorflow/core:framework",
"//tensorflow/core:framework_internal",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
],
)
tf_cc_binary(
name = "summarize_graph",
copts = tf_copts(),
linkstatic = 1,
visibility = ["//visibility:public"],
deps = [
":summarize_graph_main_lib",
],
)
tf_cc_binary(
name = "compare_graphs",
srcs = ["compare_graphs.cc"],
copts = tf_copts(),
linkstatic = 1,
visibility = ["//visibility:public"],
deps = [
":file_utils",
":transform_utils",
"//tensorflow/core:core_cpu_internal",
"//tensorflow/core:framework_internal",
"//tensorflow/core:lib",
],
)
py_library(
name = "transform_graph_py",
srcs = ["__init__.py"],
strict_deps = True,
deps = [
"//tensorflow/core:protos_all_py",
"//tensorflow/python/util:_pywrap_transform_graph",
"//tensorflow/python/util:compat",
],
)
tf_py_strict_test(
name = "transform_graph_py_test",
size = "small",
srcs = ["python/transform_graph_test.py"],
main = "python/transform_graph_test.py",
tags = ["v1only"],
deps = [
":transform_graph_py",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:tensor_util",
"//tensorflow/python/platform:client_testlib",
],
)
tensorflow/tools/graph_transforms/README.md
+1099
View File
File diff suppressed because it is too large Load Diff
tensorflow/tools/graph_transforms/__init__.py
+47
View File
@@ -0,0 +1,47 @@
# 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.
# ==============================================================================
"""Exposes the Python wrapper for graph transforms."""
# pylint: disable=unused-import,wildcard-import, line-too-long
from tensorflow.core.framework import graph_pb2
from tensorflow.python.util import compat
from tensorflow.python.util._pywrap_transform_graph import TransformGraphWithStringInputs
def TransformGraph(input_graph_def, inputs, outputs, transforms):
"""Python wrapper for the Graph Transform Tool.
Gives access to all graph transforms available through the command line tool.
See documentation at https://github.com/tensorflow/tensorflow/blob/master/tensorflow/tools/graph_transforms/README.md
for full details of the options available.
Args:
input_graph_def: GraphDef object containing a model to be transformed.
inputs: List of node names for the model inputs.
outputs: List of node names for the model outputs.
transforms: List of strings containing transform names and parameters.
Returns:
New GraphDef with transforms applied.
"""
input_graph_def_string = input_graph_def.SerializeToString()
inputs_string = compat.as_bytes(",".join(inputs))
outputs_string = compat.as_bytes(",".join(outputs))
transforms_string = compat.as_bytes(" ".join(transforms))
output_graph_def_string = TransformGraphWithStringInputs(
input_graph_def_string, inputs_string, outputs_string, transforms_string)
output_graph_def = graph_pb2.GraphDef()
output_graph_def.ParseFromString(output_graph_def_string)
return output_graph_def
tensorflow/tools/graph_transforms/add_default_attributes.cc
+40
View File
@@ -0,0 +1,40 @@
/* 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/core/framework/function.h"
#include "tensorflow/core/framework/graph_def_util.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
// Sets any parameters not specified in a node to their defaults.
absl::Status AddDefaultAttributes(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def) {
// Find all of the ops that are currently defined.
std::unique_ptr<FunctionLibraryDefinition> flib_def(
new FunctionLibraryDefinition(OpRegistry::Global(),
input_graph_def.library()));
// Works in-place, so copy over the original graph.
*output_graph_def = input_graph_def;
TF_RETURN_IF_ERROR(AddDefaultAttrsToGraphDef(output_graph_def, *flib_def, 0));
return absl::OkStatus();
}
REGISTER_GRAPH_TRANSFORM("add_default_attributes", AddDefaultAttributes);
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/add_default_attributes_test.cc
+74
View File
@@ -0,0 +1,74 @@
/* 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/cc/ops/const_op.h"
#include "tensorflow/cc/ops/image_ops.h"
#include "tensorflow/cc/ops/nn_ops.h"
#include "tensorflow/cc/ops/sendrecv_ops.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
// Declare here, so we don't need a public header.
absl::Status AddDefaultAttributes(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def);
class AddDefaultAttributesTest : public ::testing::Test {
protected:
void TestAddDefaultAttributes() {
GraphDef graph_def;
NodeDef* lrn_node1 = graph_def.add_node();
lrn_node1->set_name("lrn_node1");
lrn_node1->set_op("LRN");
NodeDef* lrn_node2 = graph_def.add_node();
lrn_node2->set_name("lrn_node2");
lrn_node2->set_op("LRN");
SetNodeAttr("depth_radius", 7, lrn_node2);
SetNodeAttr("bias", 2.0f, lrn_node2);
SetNodeAttr("alpha", 2.0f, lrn_node2);
SetNodeAttr("beta", 1.0f, lrn_node2);
GraphDef result;
TF_ASSERT_OK(AddDefaultAttributes(graph_def, {}, &result));
std::map<std::string, const NodeDef*> nodes;
MapNamesToNodes(result, &nodes);
EXPECT_EQ(5, nodes.at("lrn_node1")->attr().at("depth_radius").i());
EXPECT_NEAR(1.0f, nodes.at("lrn_node1")->attr().at("bias").f(), 1e-5f);
EXPECT_NEAR(1.0f, nodes.at("lrn_node1")->attr().at("alpha").f(), 1e-5f);
EXPECT_NEAR(0.5f, nodes.at("lrn_node1")->attr().at("beta").f(), 1e-5f);
EXPECT_EQ(7, nodes.at("lrn_node2")->attr().at("depth_radius").i());
EXPECT_NEAR(2.0f, nodes.at("lrn_node2")->attr().at("bias").f(), 1e-5f);
EXPECT_NEAR(2.0f, nodes.at("lrn_node2")->attr().at("alpha").f(), 1e-5f);
EXPECT_NEAR(1.0f, nodes.at("lrn_node2")->attr().at("beta").f(), 1e-5f);
}
};
TEST_F(AddDefaultAttributesTest, TestAddDefaultAttributes) {
TestAddDefaultAttributes();
}
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/backports.cc
+134
View File
@@ -0,0 +1,134 @@
/* 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/core/common_runtime/constant_folding.h"
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/graph/node_builder.h"
#include "tensorflow/core/graph/subgraph.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/tools/graph_transforms/fold_constants_lib.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
// Switch any ConcatV2 nodes to the v1 version, swapping the input order.
absl::Status BackportConcatV2Transform(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def) {
TF_RETURN_IF_ERROR(ReplaceMatchingOpTypes(
input_graph_def, {"ConcatV2"},
[](const NodeMatch& match, const std::set<std::string>& input_nodes,
const std::set<std::string>& output_nodes,
std::vector<NodeDef>* new_nodes) {
const NodeDef& concat_v2_node = match.node;
NodeDef concat_node = concat_v2_node;
concat_node.set_op("Concat");
// The last input is inserted at the head of the inputs, because Concat
// expects the dimension as the first input (not the last as in
// ConcatV2).
concat_node.mutable_input()->Clear();
const std::string& dim_input =
concat_v2_node.input(concat_v2_node.input_size() - 1);
concat_node.add_input(dim_input);
for (int i = 0; i < (concat_v2_node.input_size() - 1); ++i) {
concat_node.add_input(concat_v2_node.input(i));
}
// Tidx attribute must be deleted because it's not used in Concat.
concat_node.mutable_attr()->erase("Tidx");
new_nodes->push_back(concat_node);
return absl::OkStatus();
},
{true}, output_graph_def));
return absl::OkStatus();
}
REGISTER_GRAPH_TRANSFORM("backport_concatv2", BackportConcatV2Transform);
// Switch any TensorArrayV3 nodes to the v2 version, removing the second output.
absl::Status BackportTensorArrayV3Transform(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def) {
std::map<std::string, std::string> inputs_to_rename;
GraphDef replaced_graph_def;
TF_RETURN_IF_ERROR(ReplaceMatchingOpTypes(
input_graph_def, {"TensorArrayV3|TensorArrayGradV3"},
[&inputs_to_rename](const NodeMatch& match,
const std::set<std::string>& input_nodes,
const std::set<std::string>& output_nodes,
std::vector<NodeDef>* new_nodes) {
const NodeDef& tensor_array_v3_node = match.node;
// All we need to do here is rename the op type, since the attributes
// remain the same.
NodeDef tensor_array_v2_node = tensor_array_v3_node;
if (tensor_array_v3_node.op() == "TensorArrayV3") {
tensor_array_v2_node.set_op("TensorArrayV2");
} else {
tensor_array_v2_node.set_op("TensorArrayGradV2");
}
// The v3 version has a second 'flow' output that's not present in v2,
// so substitute a dummy constant instead in any places that use it.
NodeDef replacement_flow_node;
replacement_flow_node.set_op("Const");
SetNodeAttr("dtype", DT_FLOAT, &replacement_flow_node);
replacement_flow_node.set_name(tensor_array_v3_node.name() +
"/replacement_flow_node");
Tensor replacement_flow_tensor(DT_FLOAT, {});
// I'm picking an arbitrary value for the gradient flow here, for lack
// of a better alternative.
replacement_flow_tensor.flat<float>()(0) = 1.0f;
SetNodeTensorAttr<float>("value", replacement_flow_tensor,
&replacement_flow_node);
inputs_to_rename[tensor_array_v3_node.name() + ":1"] =
replacement_flow_node.name();
new_nodes->push_back(tensor_array_v2_node);
new_nodes->push_back(replacement_flow_node);
return absl::OkStatus();
},
{true}, &replaced_graph_def));
// Update the graph so that any nodes that referred to removed inputs now
// pull from the substitute constants we've added.
GraphDef renamed_graph_def;
TF_RETURN_IF_ERROR(RenameNodeInputs(replaced_graph_def, inputs_to_rename,
std::unordered_set<std::string>(),
&renamed_graph_def));
TF_RETURN_IF_ERROR(ReplaceMatchingOpTypes(
renamed_graph_def,
{"TensorArrayWriteV3|TensorArrayReadV3|TensorArrayGatherV3|"
"TensorArrayScatterV3|TensorArrayConcatV3|TensorArraySplitV3|"
"TensorArraySizeV3|TensorArrayCloseV3"},
[](const NodeMatch& match, const std::set<std::string>& input_nodes,
const std::set<std::string>& output_nodes,
std::vector<NodeDef>* new_nodes) {
const NodeDef& v3_node = match.node;
NodeDef v2_node = v3_node;
v2_node.set_op(v3_node.op().substr(0, v3_node.op().size() - 1) + "2");
new_nodes->push_back(v2_node);
return absl::OkStatus();
},
{true}, output_graph_def));
return absl::OkStatus();
}
REGISTER_GRAPH_TRANSFORM("backport_tensor_array_v3",
BackportTensorArrayV3Transform);
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/backports_test.cc
+200
View File
@@ -0,0 +1,200 @@
/* Copyright 2015 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/cc/ops/const_op.h"
#include "tensorflow/cc/ops/image_ops.h"
#include "tensorflow/cc/ops/nn_ops.h"
#include "tensorflow/cc/ops/sendrecv_ops.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/lib/strings/str_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
// Declare here, so we don't need a public header.
absl::Status BackportConcatV2Transform(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def);
absl::Status BackportTensorArrayV3Transform(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def);
class BackportConcatV2Test : public ::testing::Test {
protected:
void TestBackportConcatV2() {
GraphDef graph_def;
NodeDef* mul_node1 = graph_def.add_node();
mul_node1->set_name("mul_node1");
mul_node1->set_op("Mul");
mul_node1->add_input("add_node2");
mul_node1->add_input("add_node3");
NodeDef* add_node2 = graph_def.add_node();
add_node2->set_name("add_node2");
add_node2->set_op("Add");
add_node2->add_input("const_node1");
add_node2->add_input("const_node2");
NodeDef* add_node3 = graph_def.add_node();
add_node3->set_name("add_node3");
add_node3->set_op("Add");
add_node3->add_input("const_node1");
add_node3->add_input("const_node3");
NodeDef* const_node1 = graph_def.add_node();
const_node1->set_name("const_node1");
const_node1->set_op("Const");
NodeDef* const_node2 = graph_def.add_node();
const_node2->set_name("const_node2");
const_node2->set_op("Const");
NodeDef* const_node3 = graph_def.add_node();
const_node3->set_name("const_node3");
const_node3->set_op("Const");
NodeDef* concat_node = graph_def.add_node();
concat_node->set_name("concat_node");
concat_node->set_op("ConcatV2");
concat_node->add_input("const_node1");
concat_node->add_input("const_node2");
concat_node->add_input("const_node3");
SetNodeAttr("Tidx", DT_INT32, concat_node);
GraphDef result;
TransformFuncContext context;
context.input_names = {};
context.output_names = {"concat_node"};
TF_ASSERT_OK(BackportConcatV2Transform(graph_def, context, &result));
std::map<std::string, const NodeDef*> node_lookup;
MapNamesToNodes(result, &node_lookup);
EXPECT_EQ(1, node_lookup.count("concat_node"));
EXPECT_EQ("Concat", node_lookup.at("concat_node")->op());
EXPECT_EQ(0, node_lookup.at("concat_node")->attr().count("Tidx"));
EXPECT_EQ("const_node3", node_lookup.at("concat_node")->input(0));
EXPECT_EQ("const_node1", node_lookup.at("concat_node")->input(1));
EXPECT_EQ("const_node2", node_lookup.at("concat_node")->input(2));
EXPECT_EQ(1, node_lookup.count("const_node1"));
EXPECT_EQ("Const", node_lookup.at("const_node1")->op());
EXPECT_EQ(1, node_lookup.count("const_node2"));
EXPECT_EQ("Const", node_lookup.at("const_node2")->op());
EXPECT_EQ(1, node_lookup.count("const_node3"));
EXPECT_EQ("Const", node_lookup.at("const_node3")->op());
}
};
TEST_F(BackportConcatV2Test, TestBackportConcatV2) { TestBackportConcatV2(); }
TEST(BackportTensorArrayV3Test, TestBackportTensorArrayV3) {
GraphDef graph_def;
NodeDef* size_node = graph_def.add_node();
size_node->set_name("size_node");
size_node->set_op("Const");
Tensor size_tensor(DT_INT32, {});
size_tensor.flat<int32_t>()(0) = 1;
SetNodeTensorAttr<float>("value", size_tensor, size_node);
NodeDef* tensor_array_node = graph_def.add_node();
tensor_array_node->set_name("tensor_array_node");
tensor_array_node->set_op("TensorArrayV3");
tensor_array_node->add_input("size_node");
SetNodeAttr("dtype", DT_FLOAT, tensor_array_node);
SetNodeAttr("element_shape", TensorShape({1, 2}), tensor_array_node);
SetNodeAttr("dynamic_size", false, tensor_array_node);
SetNodeAttr("clear_after_read", true, tensor_array_node);
SetNodeAttr("tensor_array_name", "some_name", tensor_array_node);
NodeDef* handle_output_node = graph_def.add_node();
handle_output_node->set_name("handle_output_node");
handle_output_node->set_op("Identity");
handle_output_node->add_input("tensor_array_node:0");
NodeDef* flow_output_node = graph_def.add_node();
flow_output_node->set_name("flow_output_node");
flow_output_node->set_op("Identity");
flow_output_node->add_input("tensor_array_node:1");
NodeDef* tensor_array_grad_node = graph_def.add_node();
tensor_array_grad_node->set_name("tensor_array_grad_node");
tensor_array_grad_node->set_op("TensorArrayGradV3");
tensor_array_grad_node->add_input("tensor_array_node:0");
tensor_array_grad_node->add_input("tensor_array_node:1");
SetNodeAttr("source", "foo", tensor_array_grad_node);
NodeDef* grad_handle_output_node = graph_def.add_node();
grad_handle_output_node->set_name("grad_handle_output_node");
grad_handle_output_node->set_op("Identity");
grad_handle_output_node->add_input("tensor_array_grad_node:0");
NodeDef* grad_flow_output_node = graph_def.add_node();
grad_flow_output_node->set_name("grad_flow_output_node");
grad_flow_output_node->set_op("Identity");
grad_flow_output_node->add_input("tensor_array_grad_node:1");
GraphDef result;
TransformFuncContext context;
context.input_names = {};
context.output_names = {"handle_output_node", "grad_handle_output_node"};
TF_ASSERT_OK(BackportTensorArrayV3Transform(graph_def, context, &result));
std::map<std::string, const NodeDef*> node_lookup;
MapNamesToNodes(result, &node_lookup);
ASSERT_EQ(1, node_lookup.count("tensor_array_node"));
EXPECT_EQ("TensorArrayV2", node_lookup.at("tensor_array_node")->op());
EXPECT_EQ("TensorArrayGradV2",
node_lookup.at("tensor_array_grad_node")->op());
for (const NodeDef& node : result.node()) {
for (const std::string& input : node.input()) {
EXPECT_NE("tensor_array_node:1", input);
}
}
}
TEST(BackportTensorArrayV3Test, TestBackportTensorArrayV3Subtypes) {
const std::vector<std::string> v3_ops = {
"TensorArrayWriteV3", "TensorArrayReadV3", "TensorArrayGatherV3",
"TensorArrayScatterV3", "TensorArrayConcatV3", "TensorArraySplitV3",
"TensorArraySizeV3", "TensorArrayCloseV3"};
for (const std::string& v3_op : v3_ops) {
GraphDef graph_def;
NodeDef* v3_node = graph_def.add_node();
v3_node->set_name("v3_node");
v3_node->set_op(v3_op);
GraphDef result;
TransformFuncContext context;
context.input_names = {};
context.output_names = {""};
TF_ASSERT_OK(BackportTensorArrayV3Transform(graph_def, context, &result));
std::map<std::string, const NodeDef*> node_lookup;
MapNamesToNodes(result, &node_lookup);
ASSERT_EQ(1, node_lookup.count("v3_node"));
EXPECT_TRUE(absl::EndsWith(node_lookup.at("v3_node")->op(), "V2"));
}
}
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/compare_graphs.cc
+80
View File
@@ -0,0 +1,80 @@
/* 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.
==============================================================================*/
// Compares two TensorFlow graphs to see if their meaning is the same. This is a
// semantic comparison that's intended to show whether the graphs should produce
// the same results, and so ignores details like version numbers or node
// ordering that don't affect the output. To use it, run something like this:
//
// bazel build tensorflow/tools/graph_transforms:compare_graphs
// bazel-bin/tensorflow/tools/graph_transforms/compare_graphs a.pb b.pb
//
// The return value is 0 if the graphs are equal, 1 if they're different, and -1
// if there was a problem.
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/util/command_line_flags.h"
#include "tensorflow/core/util/equal_graph_def.h"
#include "tensorflow/tools/graph_transforms/file_utils.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
namespace {
int ParseFlagsAndCompareGraphs(int argc, char* argv[]) {
// We need to call this to set up global state for TensorFlow.
port::InitMain(argv[0], &argc, &argv);
if (argc != 3) {
LOG(ERROR) << "compare_graphs expects two file names as arguments";
return -1;
}
GraphDef a;
absl::Status a_load_status = LoadTextOrBinaryGraphFile(argv[1], &a);
if (!a_load_status.ok()) {
LOG(ERROR) << "Loading graph '" << argv[1] << "' failed with "
<< a_load_status.message();
return -1;
}
GraphDef b;
absl::Status b_load_status = LoadTextOrBinaryGraphFile(argv[2], &b);
if (!b_load_status.ok()) {
LOG(ERROR) << "Loading graph '" << argv[2] << "' failed with "
<< b_load_status.message();
return -1;
}
std::string diff;
if (EqualGraphDef(a, b, &diff)) {
std::cout << "Graphs are equal." << std::endl;
return 0;
} else {
std::cout << diff << std::endl;
return 1;
}
}
} // namespace
} // namespace graph_transforms
} // namespace tensorflow
int main(int argc, char* argv[]) {
return tensorflow::graph_transforms::ParseFlagsAndCompareGraphs(argc, argv);
}
tensorflow/tools/graph_transforms/file_utils.cc
+48
View File
@@ -0,0 +1,48 @@
/* 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/tools/graph_transforms/file_utils.h"
#include "tensorflow/core/platform/env.h"
namespace tensorflow {
namespace graph_transforms {
absl::Status LoadTextOrBinaryGraphFile(const std::string& file_name,
GraphDef* graph_def) {
std::string file_data;
absl::Status load_file_status =
ReadFileToString(Env::Default(), file_name, &file_data);
if (!load_file_status.ok()) {
errors::AppendToMessage(&load_file_status, " (for file ", file_name, ")");
return load_file_status;
}
// Try to load in binary format first, and then try ascii if that fails.
absl::Status load_status =
ReadBinaryProto(Env::Default(), file_name, graph_def);
if (!load_status.ok()) {
if (protobuf::TextFormat::ParseFromString(file_data, graph_def)) {
load_status = absl::OkStatus();
} else {
errors::AppendToMessage(&load_status,
" (both text and binary parsing failed for file ",
file_name, ")");
}
}
return load_status;
}
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/file_utils.h
+33
View File
@@ -0,0 +1,33 @@
/* 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_TOOLS_GRAPH_TRANSFORMS_FILE_UTILS_H_
#define TENSORFLOW_TOOLS_GRAPH_TRANSFORMS_FILE_UTILS_H_
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/lib/core/status.h"
namespace tensorflow {
namespace graph_transforms {
// First tries to load the file as a text protobuf, if that fails tries to parse
// it as a binary protobuf, and returns an error if both fail.
absl::Status LoadTextOrBinaryGraphFile(const std::string& file_name,
GraphDef* graph_def);
} // namespace graph_transforms
} // namespace tensorflow
#endif // TENSORFLOW_TOOLS_GRAPH_TRANSFORMS_FILE_UTILS_H_
tensorflow/tools/graph_transforms/file_utils_test.cc
+84
View File
@@ -0,0 +1,84 @@
/* 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/tools/graph_transforms/file_utils.h"
#include "tensorflow/cc/ops/const_op.h"
#include "tensorflow/cc/ops/image_ops.h"
#include "tensorflow/cc/ops/nn_ops.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/lib/io/path.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"
#include "tensorflow/core/util/equal_graph_def.h"
namespace tensorflow {
namespace graph_transforms {
class FileUtilsTest : public ::testing::Test {
protected:
void TestLoadTextOrBinaryGraphFile() {
using namespace ::tensorflow::ops; // NOLINT(build/namespaces)
const int width = 10;
auto root = tensorflow::Scope::NewRootScope();
Tensor a_data(DT_FLOAT, TensorShape({width}));
test::FillIota<float>(&a_data, 1.0f);
Output a_const = Const(root.WithOpName("a"), Input::Initializer(a_data));
GraphDef graph_def;
TF_ASSERT_OK(root.ToGraphDef(&graph_def));
const std::string text_file =
io::JoinPath(testing::TmpDir(), "text_graph.pbtxt");
TF_ASSERT_OK(WriteTextProto(Env::Default(), text_file, graph_def));
const std::string binary_file =
io::JoinPath(testing::TmpDir(), "binary_graph.pb");
TF_ASSERT_OK(WriteBinaryProto(Env::Default(), binary_file, graph_def));
const std::string bogus_file =
io::JoinPath(testing::TmpDir(), "bogus_graph.pb");
TF_ASSERT_OK(
WriteStringToFile(Env::Default(), bogus_file, "Not a !{ proto..."));
GraphDef text_graph_def;
TF_EXPECT_OK(LoadTextOrBinaryGraphFile(text_file, &text_graph_def));
std::string text_diff;
EXPECT_TRUE(EqualGraphDef(text_graph_def, graph_def, &text_diff))
<< text_diff;
GraphDef binary_graph_def;
TF_EXPECT_OK(LoadTextOrBinaryGraphFile(binary_file, &binary_graph_def));
std::string binary_diff;
EXPECT_TRUE(EqualGraphDef(binary_graph_def, graph_def, &binary_diff))
<< binary_diff;
GraphDef no_graph_def;
EXPECT_FALSE(
LoadTextOrBinaryGraphFile("____non_existent_file_____", &no_graph_def)
.ok());
GraphDef bogus_graph_def;
EXPECT_FALSE(LoadTextOrBinaryGraphFile(bogus_file, &bogus_graph_def).ok());
}
};
TEST_F(FileUtilsTest, TestLoadTextOrBinaryGraphFile) {
TestLoadTextOrBinaryGraphFile();
}
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/flatten_atrous.cc
+140
View File
@@ -0,0 +1,140 @@
/* 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/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/graph/node_builder.h"
#include "tensorflow/core/graph/subgraph.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
absl::Status FlattenAtrousConv(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def) {
GraphDef replaced_graph_def;
TF_RETURN_IF_ERROR(ReplaceMatchingOpTypes(
input_graph_def, // clang-format off
{"BatchToSpaceND",
{
{"Conv2D|DepthwiseConv2dNative",
{
{"SpaceToBatchND",
{
{"*"}, // Input to the flattened op.
{"*"}, // block_shape
{"*"} // paddings
}
},
{"*"} // filter
}
},
{"*"}, // block_shape
{"*"} // crops
}
}, // clang-format on
[](const NodeMatch& match, const std::set<std::string>& input_nodes,
const std::set<std::string>& output_nodes,
std::vector<NodeDef>* new_nodes) {
// Find all the nodes we expect in the subgraph.
const NodeDef& batch_to_space_node = match.node;
const NodeDef& conv_node = match.inputs[0].node;
const NodeDef& filter_node = match.inputs[0].inputs[1].node;
const NodeDef& input_node = match.inputs[0].inputs[0].inputs[0].node;
const NodeDef& space_to_batch_block_shape_node =
match.inputs[0].inputs[0].inputs[1].node;
// The atrous rate value is inferred from the block shape.
Tensor block_shape =
GetNodeTensorAttr(space_to_batch_block_shape_node, "value");
const int32_t block_height = block_shape.flat<int32_t>()(0);
const int32_t block_width = block_shape.flat<int32_t>()(1);
// Compute the upsampled filter.
const Tensor& filter = GetNodeTensorAttr(filter_node, "value");
const int32_t filter_height = filter.dim_size(0);
const int32_t filter_width = filter.dim_size(1);
const int32_t in_channels = filter.dim_size(2);
const int32_t out_channels = filter.dim_size(3);
const int32_t upsampled_filter_height =
(filter_height - 1) * block_height + 1;
const int32_t upsampled_filter_width =
(filter_width - 1) * block_width + 1;
Tensor upsampled_filter(
DT_FLOAT,
TensorShape({upsampled_filter_height, upsampled_filter_width,
in_channels, out_channels}));
auto filter_eigen = filter.tensor<float, 4>();
auto upsampled_filter_eigen = upsampled_filter.tensor<float, 4>();
upsampled_filter_eigen.setZero();
for (int h = 0; h < filter_height; ++h) {
for (int w = 0; w < filter_width; ++w) {
for (int c_in = 0; c_in < in_channels; ++c_in) {
for (int c_out = 0; c_out < out_channels; ++c_out) {
upsampled_filter_eigen(block_height * h, block_width * w, c_in,
c_out) = filter_eigen(h, w, c_in, c_out);
}
}
}
}
NodeDef upsampled_filter_node;
upsampled_filter_node.set_op("Const");
upsampled_filter_node.set_name(filter_node.name());
SetNodeAttr("dtype", DT_FLOAT, &upsampled_filter_node);
SetNodeTensorAttr<float>("value", upsampled_filter,
&upsampled_filter_node);
// Set up the new flattened version of the convolution op.
NodeDef flattened_conv_node;
flattened_conv_node.set_name(batch_to_space_node.name());
flattened_conv_node.set_op(conv_node.op());
flattened_conv_node.set_device(conv_node.device());
AddNodeInput(input_node.name(), &flattened_conv_node);
AddNodeInput(upsampled_filter_node.name(), &flattened_conv_node);
CopyNodeAttr(conv_node, "T", "T", &flattened_conv_node);
CopyNodeAttr(conv_node, "strides", "strides", &flattened_conv_node);
SetNodeAttr("padding", "SAME", &flattened_conv_node);
CopyNodeAttr(conv_node, "data_format", "data_format",
&flattened_conv_node);
if (conv_node.op() == "Conv2D") {
CopyNodeAttr(conv_node, "use_cudnn_on_gpu", "use_cudnn_on_gpu",
&flattened_conv_node);
}
new_nodes->push_back(input_node);
new_nodes->push_back(upsampled_filter_node);
new_nodes->push_back(flattened_conv_node);
return absl::OkStatus();
},
{}, &replaced_graph_def));
*output_graph_def = replaced_graph_def;
return absl::OkStatus();
}
REGISTER_GRAPH_TRANSFORM("flatten_atrous_conv", FlattenAtrousConv);
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/flatten_atrous_test.cc
+121
View File
@@ -0,0 +1,121 @@
/* 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/cc/ops/array_ops.h"
#include "tensorflow/cc/ops/const_op.h"
#include "tensorflow/cc/ops/nn_ops.h"
#include "tensorflow/cc/ops/sendrecv_ops.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
// Declare here, so we don't need a public header.
absl::Status FlattenAtrousConv(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def);
class FlattenAtrousConvTest : public ::testing::Test {
protected:
template <class TConvOp>
void TestFlattenAtrousConv() {
auto root = tensorflow::Scope::NewRootScope();
using namespace ::tensorflow::ops; // NOLINT(build/namespaces)
Tensor input_data(DT_FLOAT, TensorShape({1, 3, 3, 2}));
test::FillValues<float>(
&input_data, {.1f, .4f, .2f, .5f, .3f, .6f, -1.0f, -.4f, -.2f, -.5f,
-.3f, -.6f, .1f, .4f, .2f, .5f, .3f, .6f});
Output input_op =
Const(root.WithOpName("input_op"), Input::Initializer(input_data));
Tensor block_shape_data(DT_INT32, TensorShape({2}));
test::FillValues<int>(&block_shape_data, {2, 2});
Output block_shape_op = Const(root.WithOpName("block_shape_op"),
Input::Initializer(block_shape_data));
Tensor paddings_data(DT_INT32, TensorShape({2, 2}));
test::FillValues<int>(&paddings_data, {1, 2, 1, 2});
Output paddings_op = Const(root.WithOpName("paddings_op"),
Input::Initializer(paddings_data));
Output space_to_batch_op =
SpaceToBatchND(root.WithOpName("space_to_batch_op"), input_op,
block_shape_op, paddings_op);
Tensor weights_data(DT_FLOAT, TensorShape({2, 2, 2, 1}));
test::FillValues<float>(&weights_data,
{.1f, .2f, .3f, .4f, .1f, .2f, .3f, .4f});
Output weights_op =
Const(root.WithOpName("weights_op"), Input::Initializer(weights_data));
Output conv_op = TConvOp(root.WithOpName("conv_op"), space_to_batch_op,
weights_op, {1, 1, 1, 1}, "VALID");
Tensor crops_data(DT_INT32, TensorShape({2, 2}));
test::FillValues<int>(&crops_data, {0, 1, 0, 1});
Output crops_op =
Const(root.WithOpName("crops_op"), Input::Initializer(crops_data));
Output batch_to_space_op = BatchToSpaceND(
root.WithOpName("output"), conv_op, block_shape_op, crops_op);
GraphDef original_graph_def;
TF_ASSERT_OK(root.ToGraphDef(&original_graph_def));
std::unique_ptr<Session> original_session(NewSession(SessionOptions()));
TF_ASSERT_OK(original_session->Create(original_graph_def));
std::vector<Tensor> original_outputs;
TF_ASSERT_OK(original_session->Run({}, {"output"}, {}, &original_outputs));
GraphDef modified_graph_def;
TF_ASSERT_OK(FlattenAtrousConv(original_graph_def, {{}, {"output"}},
&modified_graph_def));
std::unique_ptr<Session> modified_session(NewSession(SessionOptions()));
TF_ASSERT_OK(modified_session->Create(modified_graph_def));
std::vector<Tensor> modified_outputs;
TF_ASSERT_OK(modified_session->Run({}, {"output"}, {}, &modified_outputs));
EXPECT_EQ(3, modified_graph_def.node_size());
EXPECT_EQ("input_op", modified_graph_def.node(0).name());
EXPECT_EQ("weights_op", modified_graph_def.node(1).name());
EXPECT_EQ("output", modified_graph_def.node(2).name());
EXPECT_EQ("Const", modified_graph_def.node(0).op());
EXPECT_EQ("Const", modified_graph_def.node(1).op());
EXPECT_EQ(conv_op.node()->type_string(), modified_graph_def.node(2).op());
test::ExpectTensorNear<float>(original_outputs[0], modified_outputs[0],
1e-6);
}
};
TEST_F(FlattenAtrousConvTest, TestFlattenAtrousConv2D) {
TestFlattenAtrousConv<::tensorflow::ops::Conv2D>();
}
TEST_F(FlattenAtrousConvTest, TestFlattenAtrousDepthwiseConv2dNative) {
TestFlattenAtrousConv<::tensorflow::ops::DepthwiseConv2dNative>();
}
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/fold_batch_norms.cc
+125
View File
@@ -0,0 +1,125 @@
/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/core/common_runtime/constant_folding.h"
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/graph/node_builder.h"
#include "tensorflow/core/graph/subgraph.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/tools/graph_transforms/fold_constants_lib.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
// Converts Conv2D or MatMul ops followed by column-wise Muls into equivalent
// ops with the Mul baked into the convolution weights, to save computation
// during inference.
absl::Status FoldBatchNorms(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def) {
GraphDef replaced_graph_def;
TF_RETURN_IF_ERROR(ReplaceMatchingOpTypes(
input_graph_def, // clang-format off
{"Mul", // mul_node
{
{"Conv2D|MatMul|DepthwiseConv2dNative", // conv_node
{
{"*"}, // input_node
{"Const"}, // weights_node
}
},
{"Const"}, // mul_values_node
}
}, // clang-format on
[](const NodeMatch& match, const std::set<std::string>& input_nodes,
const std::set<std::string>& output_nodes,
std::vector<NodeDef>* new_nodes) {
// Find all the nodes we expect in the subgraph.
const NodeDef& mul_node = match.node;
const NodeDef& conv_node = match.inputs[0].node;
const NodeDef& input_node = match.inputs[0].inputs[0].node;
const NodeDef& weights_node = match.inputs[0].inputs[1].node;
const NodeDef& mul_values_node = match.inputs[1].node;
// Check that nodes that we use are not used somewhere else.
for (const auto& node : {conv_node, weights_node, mul_values_node}) {
if (output_nodes.count(node.name())) {
// Return original nodes.
new_nodes->insert(new_nodes->end(),
{mul_node, conv_node, input_node, weights_node,
mul_values_node});
return absl::OkStatus();
}
}
Tensor weights = GetNodeTensorAttr(weights_node, "value");
Tensor mul_values = GetNodeTensorAttr(mul_values_node, "value");
// Make sure all the inputs really are vectors, with as many entries as
// there are columns in the weights.
int64_t weights_cols;
if (conv_node.op() == "Conv2D") {
weights_cols = weights.shape().dim_size(3);
} else if (conv_node.op() == "DepthwiseConv2dNative") {
weights_cols =
weights.shape().dim_size(2) * weights.shape().dim_size(3);
} else {
weights_cols = weights.shape().dim_size(1);
}
if ((mul_values.shape().dims() != 1) ||
(mul_values.shape().dim_size(0) != weights_cols)) {
return absl::InvalidArgumentError(
absl::StrCat("Mul constant input to batch norm has bad shape: ",
mul_values.shape().DebugString()));
}
// Multiply the original weights by the scale vector.
auto weights_vector = weights.flat<float>();
Tensor scaled_weights(DT_FLOAT, weights.shape());
auto scaled_weights_vector = scaled_weights.flat<float>();
for (int64_t row = 0; row < weights_vector.dimension(0); ++row) {
scaled_weights_vector(row) =
weights_vector(row) *
mul_values.flat<float>()(row % weights_cols);
}
// Construct the new nodes.
NodeDef scaled_weights_node;
scaled_weights_node.set_op("Const");
scaled_weights_node.set_name(weights_node.name());
SetNodeAttr("dtype", DT_FLOAT, &scaled_weights_node);
SetNodeTensorAttr<float>("value", scaled_weights, &scaled_weights_node);
new_nodes->push_back(scaled_weights_node);
new_nodes->push_back(input_node);
NodeDef new_conv_node;
new_conv_node = conv_node;
new_conv_node.set_name(mul_node.name());
new_nodes->push_back(new_conv_node);
return absl::OkStatus();
},
{}, &replaced_graph_def));
*output_graph_def = replaced_graph_def;
return absl::OkStatus();
}
REGISTER_GRAPH_TRANSFORM("fold_batch_norms", FoldBatchNorms);
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/fold_batch_norms_test.cc
+261
View File
@@ -0,0 +1,261 @@
/* Copyright 2015 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/cc/ops/const_op.h"
#include "tensorflow/cc/ops/image_ops.h"
#include "tensorflow/cc/ops/math_ops.h"
#include "tensorflow/cc/ops/nn_ops.h"
#include "tensorflow/cc/ops/sendrecv_ops.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
// Declare here, so we don't need a public header.
absl::Status FoldBatchNorms(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def);
class FoldBatchNormsTest : public ::testing::Test {
protected:
void TestFoldBatchNormsConv2D() {
auto root = tensorflow::Scope::NewRootScope();
using namespace ::tensorflow::ops; // NOLINT(build/namespaces)
Tensor input_data(DT_FLOAT, TensorShape({1, 1, 6, 2}));
test::FillValues<float>(
&input_data, {1.0f, 4.0f, 2.0f, 5.0f, 3.0f, 6.0f, -1.0f, -4.0f, -2.0f,
-5.0f, -3.0f, -6.0f});
Output input_op =
Const(root.WithOpName("input_op"), Input::Initializer(input_data));
Tensor weights_data(DT_FLOAT, TensorShape({1, 2, 2, 2}));
test::FillValues<float>(&weights_data,
{1.0f, 2.0f, 3.0f, 4.0f, 0.1f, 0.2f, 0.3f, 0.4f});
Output weights_op =
Const(root.WithOpName("weights_op"), Input::Initializer(weights_data));
Output conv_op = Conv2D(root.WithOpName("conv_op"), input_op, weights_op,
{1, 1, 1, 1}, "VALID");
Tensor mul_values_data(DT_FLOAT, TensorShape({2}));
test::FillValues<float>(&mul_values_data, {2.0f, 3.0f});
Output mul_values_op = Const(root.WithOpName("mul_values"),
Input::Initializer(mul_values_data));
Output mul_op = Mul(root.WithOpName("output"), conv_op, mul_values_op);
GraphDef original_graph_def;
TF_ASSERT_OK(root.ToGraphDef(&original_graph_def));
std::unique_ptr<Session> original_session(NewSession(SessionOptions()));
TF_ASSERT_OK(original_session->Create(original_graph_def));
std::vector<Tensor> original_outputs;
TF_ASSERT_OK(original_session->Run({}, {"output"}, {}, &original_outputs));
GraphDef fused_graph_def;
TF_ASSERT_OK(
FoldBatchNorms(original_graph_def, {{}, {"output"}}, &fused_graph_def));
std::unique_ptr<Session> fused_session(NewSession(SessionOptions()));
TF_ASSERT_OK(fused_session->Create(fused_graph_def));
std::vector<Tensor> fused_outputs;
TF_ASSERT_OK(fused_session->Run({}, {"output"}, {}, &fused_outputs));
test::ExpectTensorNear<float>(original_outputs[0], fused_outputs[0], 1e-5);
for (const NodeDef& node : fused_graph_def.node()) {
EXPECT_NE("Mul", node.op());
}
}
void TestFoldBatchNormsDepthwiseConv2dNative() {
auto root = tensorflow::Scope::NewRootScope();
using namespace ::tensorflow::ops; // NOLINT(build/namespaces)
Tensor input_data(DT_FLOAT, TensorShape({1, 1, 6, 2}));
test::FillValues<float>(
&input_data, {1.0f, 4.0f, 2.0f, 5.0f, 3.0f, 6.0f, -1.0f, -4.0f, -2.0f,
-5.0f, -3.0f, -6.0f});
Output input_op =
Const(root.WithOpName("input_op"), Input::Initializer(input_data));
Tensor weights_data(DT_FLOAT, TensorShape({1, 2, 2, 2}));
test::FillValues<float>(&weights_data,
{1.0f, 2.0f, 3.0f, 4.0f, 0.1f, 0.2f, 0.3f, 0.4f});
Output weights_op =
Const(root.WithOpName("weights_op"), Input::Initializer(weights_data));
Output conv_op = DepthwiseConv2dNative(root.WithOpName("conv_op"), input_op,
weights_op, {1, 1, 1, 1}, "VALID");
Tensor mul_values_data(DT_FLOAT, TensorShape({4}));
test::FillValues<float>(&mul_values_data, {2.0f, 3.0f, 4.0f, 5.0f});
Output mul_values_op = Const(root.WithOpName("mul_values"),
Input::Initializer(mul_values_data));
Output mul_op = Mul(root.WithOpName("output"), conv_op, mul_values_op);
GraphDef original_graph_def;
TF_ASSERT_OK(root.ToGraphDef(&original_graph_def));
std::unique_ptr<Session> original_session(NewSession(SessionOptions()));
TF_ASSERT_OK(original_session->Create(original_graph_def));
std::vector<Tensor> original_outputs;
TF_ASSERT_OK(original_session->Run({}, {"output"}, {}, &original_outputs));
GraphDef fused_graph_def;
TF_ASSERT_OK(
FoldBatchNorms(original_graph_def, {{}, {"output"}}, &fused_graph_def));
std::unique_ptr<Session> fused_session(NewSession(SessionOptions()));
TF_ASSERT_OK(fused_session->Create(fused_graph_def));
std::vector<Tensor> fused_outputs;
TF_ASSERT_OK(fused_session->Run({}, {"output"}, {}, &fused_outputs));
test::ExpectTensorNear<float>(original_outputs[0], fused_outputs[0], 1e-5);
for (const NodeDef& node : fused_graph_def.node()) {
EXPECT_NE("Mul", node.op());
}
}
void TestFoldBatchNormsConv2DShared() {
auto root = tensorflow::Scope::NewRootScope();
using namespace ::tensorflow::ops; // NOLINT(build/namespaces)
Tensor input_data(DT_FLOAT, TensorShape({1, 1, 6, 2}));
test::FillValues<float>(
&input_data, {1.0f, 4.0f, 2.0f, 5.0f, 3.0f, 6.0f, -1.0f, -4.0f, -2.0f,
-5.0f, -3.0f, -6.0f});
Output input_op =
Const(root.WithOpName("input_op"), Input::Initializer(input_data));
Tensor weights_data(DT_FLOAT, TensorShape({1, 2, 2, 2}));
test::FillValues<float>(&weights_data,
{1.0f, 2.0f, 3.0f, 4.0f, 0.1f, 0.2f, 0.3f, 0.4f});
Output weights_op =
Const(root.WithOpName("weights_op"), Input::Initializer(weights_data));
Output conv_op = Conv2D(root.WithOpName("conv_op"), input_op, weights_op,
{1, 1, 1, 1}, "VALID");
Tensor mul_values_data(DT_FLOAT, TensorShape({2}));
test::FillValues<float>(&mul_values_data, {2.0f, 3.0f});
Output mul_values_op = Const(root.WithOpName("mul_values"),
Input::Initializer(mul_values_data));
Output mul_op = Mul(root.WithOpName("output"), conv_op, mul_values_op);
Tensor mul_values_data_2(DT_FLOAT, TensorShape({2}));
test::FillValues<float>(&mul_values_data_2, {1.0f, 2.0f});
Output mul_values_op_2 = Const(root.WithOpName("mul_values_2"),
Input::Initializer(mul_values_data));
Output mul_op_2 =
Mul(root.WithOpName("output_2"), conv_op, mul_values_op_2);
GraphDef original_graph_def;
TF_ASSERT_OK(root.ToGraphDef(&original_graph_def));
std::unique_ptr<Session> original_session(NewSession(SessionOptions()));
TF_ASSERT_OK(original_session->Create(original_graph_def));
std::vector<Tensor> original_outputs;
TF_ASSERT_OK(original_session->Run({}, {"output", "output_2"}, {},
&original_outputs));
GraphDef fused_graph_def;
TF_ASSERT_OK(FoldBatchNorms(
original_graph_def, {{}, {"output", "output_2"}}, &fused_graph_def));
std::unique_ptr<Session> fused_session(NewSession(SessionOptions()));
TF_ASSERT_OK(fused_session->Create(fused_graph_def));
std::vector<Tensor> fused_outputs;
TF_ASSERT_OK(
fused_session->Run({}, {"output", "output_2"}, {}, &fused_outputs));
test::ExpectTensorNear<float>(original_outputs[0], fused_outputs[0], 1e-5);
test::ExpectTensorNear<float>(original_outputs[1], fused_outputs[1], 1e-5);
}
void TestFoldBatchNormsMatMul() {
auto root = tensorflow::Scope::NewRootScope();
using namespace ::tensorflow::ops; // NOLINT(build/namespaces)
Tensor input_data(DT_FLOAT, TensorShape({6, 2}));
test::FillValues<float>(
&input_data, {1.0f, 4.0f, 2.0f, 5.0f, 3.0f, 6.0f, -1.0f, -4.0f, -2.0f,
-5.0f, -3.0f, -6.0f});
Output input_op =
Const(root.WithOpName("input_op"), Input::Initializer(input_data));
Tensor weights_data(DT_FLOAT, TensorShape({2, 2}));
test::FillValues<float>(&weights_data, {1.0f, 2.0f, 0.3f, 0.4f});
Output weights_op =
Const(root.WithOpName("weights_op"), Input::Initializer(weights_data));
Output matmul_op =
MatMul(root.WithOpName("matmul_op"), input_op, weights_op);
Tensor mul_values_data(DT_FLOAT, TensorShape({2}));
test::FillValues<float>(&mul_values_data, {2.0f, 3.0f});
Output mul_values_op = Const(root.WithOpName("mul_values"),
Input::Initializer(mul_values_data));
Output mul_op = Mul(root.WithOpName("output"), matmul_op, mul_values_op);
GraphDef original_graph_def;
TF_ASSERT_OK(root.ToGraphDef(&original_graph_def));
std::unique_ptr<Session> original_session(NewSession(SessionOptions()));
TF_ASSERT_OK(original_session->Create(original_graph_def));
std::vector<Tensor> original_outputs;
TF_ASSERT_OK(original_session->Run({}, {"output"}, {}, &original_outputs));
GraphDef fused_graph_def;
TF_ASSERT_OK(
FoldBatchNorms(original_graph_def, {{}, {"output"}}, &fused_graph_def));
std::unique_ptr<Session> fused_session(NewSession(SessionOptions()));
TF_ASSERT_OK(fused_session->Create(fused_graph_def));
std::vector<Tensor> fused_outputs;
TF_ASSERT_OK(fused_session->Run({}, {"output"}, {}, &fused_outputs));
test::ExpectTensorNear<float>(original_outputs[0], fused_outputs[0], 1e-5);
for (const NodeDef& node : fused_graph_def.node()) {
EXPECT_NE("Mul", node.op());
}
}
};
TEST_F(FoldBatchNormsTest, TestFoldBatchNormsConv2D) {
TestFoldBatchNormsConv2D();
}
TEST_F(FoldBatchNormsTest, TestFoldBatchNormsMatMul) {
TestFoldBatchNormsMatMul();
}
TEST_F(FoldBatchNormsTest, TestFoldBatchNormsDepthwiseConv2dNative) {
TestFoldBatchNormsDepthwiseConv2dNative();
}
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/fold_constants_lib.cc
+341
View File
@@ -0,0 +1,341 @@
/* Copyright 2016 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/tools/graph_transforms/fold_constants_lib.h"
#include <algorithm>
#include <iterator>
#include <map>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>
#include "tensorflow/core/common_runtime/constant_folding.h"
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/common_runtime/shape_refiner.h"
#include "tensorflow/core/graph/node_builder.h"
#include "tensorflow/core/graph/subgraph.h"
#include "tensorflow/core/lib/core/stringpiece.h"
#include "tensorflow/core/lib/strings/numbers.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
namespace {
using StringPieceSet = std::unordered_set<absl::string_view, StringPieceHasher>;
template <typename T>
using StringPieceMap =
std::unordered_map<absl::string_view, T, StringPieceHasher>;
} // namespace
absl::Status ReplaceSendRecvs(const GraphDef& original_graph_def,
const GraphDef& rewritten_graph_def,
const std::vector<std::string>& inputs,
const std::vector<std::string>& outputs,
GraphDef* output_graph_def) {
// recv_node_names serves as a string storage for recv node names.
std::vector<std::string> recv_node_names(inputs.size());
StringPieceMap<TensorId> recv_node_map;
StringPieceSet input_nodes;
for (int i = 0; i < inputs.size(); ++i) {
// RewriteGraphForExecution adds a recv node for each input edge. We assume
// here that adding such recv node did not fail. For example, the original
// graph did not already have a node with the name for the new added recv
// node.
TensorId id = ParseTensorName(inputs[i]);
input_nodes.insert(id.first);
std::string& recv_node_name = recv_node_names[i];
recv_node_name = absl::StrCat("_recv_", id.first, "_", id.second);
recv_node_map.emplace(recv_node_name, id);
}
StringPieceMap<const NodeDef*> original_map;
for (const NodeDef& node : original_graph_def.node()) {
original_map.emplace(node.name(), &node);
}
for (const NodeDef& node : rewritten_graph_def.node()) {
if ((node.op() == "_Send") || (node.op() == "_Recv")) {
// If the op is a Send or Recv that wasn't in the original, skip it.
if (original_map.count(node.name()) == 0) {
continue;
}
}
NodeDef* new_node = output_graph_def->add_node();
new_node->MergeFrom(node);
for (int i = 0; i < new_node->input_size(); ++i) {
std::string& input = *new_node->mutable_input(i);
TensorId id = ParseTensorName(input);
const auto iter = recv_node_map.find(id.first);
if (iter != recv_node_map.end()) {
// The node being substituted is a Recv node, and it has only one
// output. If this input is not a control input, then replace the input
// with the mapped value. Otherwise, replace the node name only.
if (id.second != Graph::kControlSlot) {
CHECK_EQ(id.second, 0);
input = iter->second.ToString();
} else {
id.first = iter->second.first;
input = id.ToString();
}
}
}
// RewriteGraphForExecution() did not remove this input node. Remove this
// node name from input_nodes so that a duplicate does not get added to the
// output_graph_def.
auto iter = input_nodes.find(new_node->name());
if (iter != input_nodes.end()) {
input_nodes.erase(iter);
}
}
// Some input nodes are removed in rewrite_graph_def. Add those nodes to
// output_graph_def.
for (absl::string_view name : input_nodes) {
const NodeDef& removed_node = *CHECK_NOTNULL(original_map[name]);
output_graph_def->add_node()->MergeFrom(removed_node);
}
return absl::OkStatus();
}
absl::Status RewriteInputsAsPlaceholders(const TransformFuncContext& context,
GraphDef* graph_def) {
std::unordered_set<std::string> input_names;
for (const std::string& input_name : context.input_names) {
input_names.emplace(ParseTensorName(input_name).first);
}
for (NodeDef& node : *graph_def->mutable_node()) {
if (input_names.find(node.name()) == input_names.end()) {
continue;
}
if (node.op() == "PlaceholderWithDefault") {
node.set_op("Placeholder");
node.clear_input();
} else if (node.op() != "Placeholder") {
return absl::InvalidArgumentError(absl::StrCat(
"Input '", node.name(),
"' was expected to be a Placeholder or PlaceholderWithDefault op, "
"but was ",
node.op()));
}
}
return absl::OkStatus();
}
absl::Status RemoveUnusedNodes(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def) {
StringPieceMap<const NodeDef*> node_map;
for (const NodeDef& node : input_graph_def.node()) {
node_map.emplace(node.name(), &node);
}
std::unordered_set<TensorId, TensorId::Hasher> input_names;
for (const std::string& input : context.input_names) {
input_names.insert(ParseTensorName(input));
}
StringPieceSet used_nodes;
StringPieceSet current_nodes;
for (const std::string& name : context.output_names) {
TensorId id = ParseTensorName(name);
used_nodes.insert(id.first);
current_nodes.insert(id.first);
}
while (!current_nodes.empty()) {
StringPieceSet next_nodes;
for (absl::string_view node_name : current_nodes) {
if (node_map.count(node_name) == 0) {
LOG(ERROR) << "Bad graph structure, no node named '" << node_name
<< "' found for input lookup";
return absl::InvalidArgumentError(
absl::StrCat("Bad graph structure, no node named '", node_name,
"' found for input lookup"));
}
const NodeDef& node = *(node_map[node_name]);
for (const std::string& input : node.input()) {
TensorId id = ParseTensorName(input);
if (input_names.count(id) > 0) {
continue;
}
if (used_nodes.insert(id.first).second) {
next_nodes.insert(id.first);
}
}
}
current_nodes.swap(next_nodes);
}
for (const TensorId& id : input_names) {
used_nodes.insert(id.first);
}
FilterGraphDef(
input_graph_def,
[&](const NodeDef& node) { return used_nodes.count(node.name()) > 0; },
output_graph_def);
TF_RETURN_IF_ERROR(RewriteInputsAsPlaceholders(context, output_graph_def));
return absl::OkStatus();
}
// Converts a shape inference handle to a PartialTensorShape.
absl::Status ShapeHandleToTensorShape(
const shape_inference::ShapeHandle& handle,
shape_inference::InferenceContext* context, PartialTensorShape* shape) {
// The default is already unknown.
if (!context->RankKnown(handle)) return absl::OkStatus();
std::vector<int64_t> dims(context->Rank(handle));
for (int32_t i = 0; i < dims.size(); ++i) {
dims[i] = context->Value(context->Dim(handle, i));
}
return PartialTensorShape::MakePartialShape(dims.data(), dims.size(), shape);
}
// Converts any sub-graphs that can be resolved into constant expressions into
// single Const ops.
absl::Status FoldConstants(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def) {
Graph input_graph(OpRegistry::Global());
TF_RETURN_IF_ERROR(input_graph.AddFunctionLibrary(input_graph_def.library()));
ShapeRefiner shape_refiner(input_graph.versions(), input_graph.op_registry());
shape_refiner.set_require_shape_inference_fns(false);
shape_refiner.set_disable_constant_propagation(false);
shape_refiner.set_function_library_for_shape_inference(
&input_graph.flib_def());
bool clear_output_shapes;
TF_RETURN_IF_ERROR(context.GetOneBoolParameter("clear_output_shapes", true,
&clear_output_shapes));
if (clear_output_shapes) {
// Some older GraphDefs have saved _output_shapes attributes which are out
// of date and cause import errors, so clean them up first.
GraphDef cleaned_graph_def;
RemoveAttributes(input_graph_def, {"_output_shapes"}, &cleaned_graph_def);
TF_RETURN_IF_ERROR(
ImportGraphDef({}, cleaned_graph_def, &input_graph, &shape_refiner));
} else {
TF_RETURN_IF_ERROR(
ImportGraphDef({}, input_graph_def, &input_graph, &shape_refiner));
}
// Sorted array of input names as lookup table.
std::vector<TensorId> input_names;
input_names.reserve(context.input_names.size());
std::transform(context.input_names.begin(), context.input_names.end(),
std::back_inserter(input_names),
[](const std::string& name) { return ParseTensorName(name); });
const auto compare = [](TensorId lhs, TensorId rhs) {
return lhs.first < rhs.first;
};
std::sort(input_names.begin(), input_names.end(), compare);
// Set statically inferred shapes.
std::unordered_map<std::string, std::vector<PartialTensorShape>> shape_map;
for (const Node* const node : input_graph.nodes()) {
auto ctx = shape_refiner.GetContext(node);
if (ctx == nullptr) {
continue;
}
std::vector<PartialTensorShape>& partial_shapes = shape_map[node->name()];
if (ctx->num_outputs() <= 0) continue;
partial_shapes.resize(ctx->num_outputs());
// Check all outputs.
for (const Edge* out_edge : node->out_edges()) {
if (out_edge->IsControlEdge()) continue;
const int output_idx = out_edge->src_output();
TF_RETURN_IF_ERROR(ShapeHandleToTensorShape(ctx->output(output_idx), ctx,
&partial_shapes[output_idx]));
}
// RewriteGraphForExecution() will add a Recv node for each input. Shape
// refiner does not include shape information of these Recv nodes. Therefore
// we add entries for Recv nodes here.
const auto pair = std::equal_range(input_names.begin(), input_names.end(),
TensorId{node->name(), 0}, compare);
for (auto it = pair.first; it != pair.second; ++it) {
const std::string recv_name =
absl::StrCat("_recv_", it->first, "_", it->second);
auto& recv_partial_shapes = shape_map[recv_name];
// For whatever reason (for example, name collision) if the map entry was
// already there, then do nothing.
if (recv_partial_shapes.empty()) {
recv_partial_shapes.push_back(partial_shapes[it->second]);
}
}
}
subgraph::RewriteGraphMetadata unused_metadata;
TF_RETURN_IF_ERROR(subgraph::RewriteGraphForExecution(
&input_graph, context.input_names, context.output_names, {}, {},
false /* use_function_convention */, &unused_metadata));
ConstantFoldingOptions cf_opts;
cf_opts.shape_map = &shape_map;
// Exclude specified nodes from constant folding.
std::set<std::string> excluded_ops, excluded_nodes;
if (context.params.count("exclude_op") > 0) {
const auto& ops = context.params.at("exclude_op");
excluded_ops = std::set<std::string>(ops.begin(), ops.end());
}
if (context.params.count("exclude_node") > 0) {
const auto& nodes = context.params.at("exclude_node");
excluded_nodes = std::set<std::string>(nodes.begin(), nodes.end());
}
if (!excluded_ops.empty() || !excluded_nodes.empty()) {
cf_opts.consider = [excluded_ops, excluded_nodes](const Node* n) {
return excluded_ops.find(n->op_def().name()) == excluded_ops.end() &&
excluded_nodes.find(n->name()) == excluded_nodes.end();
};
}
TF_RETURN_IF_ERROR(context.GetOneInt64Parameter(
"max_constant_size_in_bytes", cf_opts.max_constant_size_in_bytes,
&cf_opts.max_constant_size_in_bytes));
// Constant folding.
bool was_mutated;
TF_RETURN_IF_ERROR(ConstantFold(cf_opts, nullptr, Env::Default(), nullptr,
&input_graph, &was_mutated));
GraphDef folded_graph_def;
input_graph.ToGraphDef(&folded_graph_def);
GraphDef send_recvs_replaced;
TF_RETURN_IF_ERROR(ReplaceSendRecvs(input_graph_def, folded_graph_def,
context.input_names, context.output_names,
&send_recvs_replaced));
TF_RETURN_IF_ERROR(
RemoveUnusedNodes(send_recvs_replaced, context, output_graph_def));
return absl::OkStatus();
}
REGISTER_GRAPH_TRANSFORM("fold_constants", FoldConstants);
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/fold_constants_lib.h
+43
View File
@@ -0,0 +1,43 @@
/* Copyright 2015 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_TOOLS_GRAPH_TRANSFORMS_FOLD_CONSTANTS_LIB_H_
#define TENSORFLOW_TOOLS_GRAPH_TRANSFORMS_FOLD_CONSTANTS_LIB_H_
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
// Finds sub-graphs that evaluate to constant expressions, and replaces them
// with Const nodes, to simplify the graph. The inputs and outputs arguments are
// the names of all the nodes that data is fed into, or read out of, when the
// graph is actually run.
absl::Status FoldConstants(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def);
// Analyzes which nodes are used for the given set of inputs and outputs, and
// returns a copy of the graph with any that aren't used removed.
absl::Status RemoveUnusedNodes(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def);
} // namespace graph_transforms
} // namespace tensorflow
#endif // TENSORFLOW_TOOLS_GRAPH_TRANSFORMS_FOLD_CONSTANTS_LIB_H_
tensorflow/tools/graph_transforms/fold_constants_test.cc
+397
View File
@@ -0,0 +1,397 @@
/* Copyright 2015 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 <utility>
#include "tensorflow/cc/ops/const_op.h"
#include "tensorflow/cc/ops/image_ops.h"
#include "tensorflow/cc/ops/nn_ops.h"
#include "tensorflow/cc/ops/sendrecv_ops.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/tensor_shape.pb.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/lib/strings/str_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/tools/graph_transforms/fold_constants_lib.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
// Declaring this here so it doesn't need to be in the public header.
absl::Status ReplaceSendRecvs(const GraphDef& original_graph_def,
const GraphDef& rewritten_graph_def,
const std::vector<std::string>& inputs,
const std::vector<std::string>& outputs,
GraphDef* output_graph_def);
class ConstantFoldingTest : public ::testing::Test {
protected:
void TestSimpleAdd() {
auto root = tensorflow::Scope::NewRootScope();
using namespace ::tensorflow::ops; // NOLINT(build/namespaces)
const int width = 100;
Tensor a_data(DT_FLOAT, TensorShape({width}));
test::FillIota<float>(&a_data, 1.0f);
Output a_const =
Const(root.WithOpName("a_expect_removed"), Input::Initializer(a_data));
Tensor b_data(DT_FLOAT, TensorShape({width}));
test::FillIota<float>(&b_data, 1.0f);
Output b_const =
Const(root.WithOpName("b_expect_removed"), Input::Initializer(b_data));
Output add = Add(root.WithOpName("add_expect_removed"), a_const, b_const);
Output placeholder =
Placeholder(root.WithOpName("placeholder_expect_remains"), DT_FLOAT);
Output mul =
Mul(root.WithOpName("output_expect_remains"), add, placeholder);
GraphDef graph_def;
TF_ASSERT_OK(root.ToGraphDef(&graph_def));
Tensor placeholder_tensor(DT_FLOAT, TensorShape({width}));
test::FillIota<float>(&placeholder_tensor, 1.0f);
TestConstantFolding(graph_def,
{{"placeholder_expect_remains", placeholder_tensor}},
{}, {"output_expect_remains"}, {});
TestConstantFolding(graph_def,
{{"placeholder_expect_remains:0", placeholder_tensor}},
{}, {"output_expect_remains:0"}, {});
}
void TestOpExclusionAdd() {
auto root = tensorflow::Scope::NewRootScope();
using namespace ::tensorflow::ops; // NOLINT(build/namespaces)
const int width = 100;
Tensor a_data(DT_FLOAT, TensorShape({width}));
test::FillIota<float>(&a_data, 1.0f);
Output a_const =
Const(root.WithOpName("a_expect_remains"), Input::Initializer(a_data));
Tensor b_data(DT_FLOAT, TensorShape({width}));
test::FillIota<float>(&b_data, 1.0f);
Output b_const =
Const(root.WithOpName("b_expect_remains"), Input::Initializer(b_data));
Output add = Add(root.WithOpName("add_expect_remains"), a_const, b_const);
Output placeholder =
Placeholder(root.WithOpName("placeholder_expect_remains"), DT_FLOAT);
Output mul =
Mul(root.WithOpName("output_expect_remains"), add, placeholder);
GraphDef graph_def;
TF_ASSERT_OK(root.ToGraphDef(&graph_def));
Tensor placeholder_tensor(DT_FLOAT, TensorShape({width}));
test::FillIota<float>(&placeholder_tensor, 1.0f);
TestConstantFolding(graph_def,
{{"placeholder_expect_remains", placeholder_tensor}},
{"Add"}, {"output_expect_remains"}, {});
}
void TestShapePropagation() {
auto root = tensorflow::Scope::NewRootScope();
using namespace ::tensorflow::ops; // NOLINT(build/namespaces)
Output placeholder =
Placeholder(root.WithOpName("placeholder_expect_remains"), DT_FLOAT);
Output a_const =
Const(root.WithOpName("a_expect_removed"),
Input::Initializer({1, 1, 1}, TensorShape({1, 1, 3})));
Output shape = Shape(root.WithOpName("shape_expect_removed"), a_const);
Output cast = Cast(root.WithOpName("cast_expect_removed"), shape, DT_FLOAT);
Output mul =
Mul(root.WithOpName("output_expect_remains"), cast, placeholder);
GraphDef graph_def;
TF_ASSERT_OK(root.ToGraphDef(&graph_def));
Tensor placeholder_tensor(DT_FLOAT, TensorShape({3}));
test::FillIota<float>(&placeholder_tensor, 1.0);
TestConstantFolding(graph_def,
{{"placeholder_expect_remains", placeholder_tensor}},
{}, {"output_expect_remains"}, {});
}
void TestPreserveOutputShapes() {
auto root = tensorflow::Scope::NewRootScope();
using namespace ::tensorflow::ops; // NOLINT(build/namespaces)
tensorflow::AttrValue shape_attr;
auto* shape_proto = shape_attr.mutable_list()->add_shape();
shape_proto->add_dim()->set_size(1);
shape_proto->add_dim()->set_size(1);
shape_proto->add_dim()->set_size(3);
Output placeholder =
Placeholder(root.WithOpName("placeholder_expect_remains"), DT_FLOAT);
placeholder.node()->AddAttr("_output_shapes", shape_attr);
Output shape = Shape(root.WithOpName("shape_expect_removed"), placeholder);
Output cast = Cast(root.WithOpName("cast_expect_removed"), shape, DT_FLOAT);
Output mul =
Mul(root.WithOpName("output_expect_remains"), cast, placeholder);
GraphDef graph_def;
TF_ASSERT_OK(root.ToGraphDef(&graph_def));
Tensor placeholder_tensor(DT_FLOAT, TensorShape({1, 1, 3}));
test::FillIota<float>(&placeholder_tensor, 1.0);
graph_transforms::TransformFuncContext context;
context.params["clear_output_shapes"] = {"false"};
TestConstantFolding(graph_def,
{{"placeholder_expect_remains", placeholder_tensor}},
{}, {"output_expect_remains"}, context);
}
void TestConstantFolding(const GraphDef& graph_def,
std::vector<std::pair<std::string, Tensor> > inputs,
std::vector<std::string> excluded_ops,
const std::vector<std::string>& outputs,
graph_transforms::TransformFuncContext context) {
std::unique_ptr<tensorflow::Session> unfolded_session(
tensorflow::NewSession(tensorflow::SessionOptions()));
TF_ASSERT_OK(unfolded_session->Create(graph_def));
std::vector<Tensor> unfolded_tensors;
TF_ASSERT_OK(unfolded_session->Run(inputs, outputs, {}, &unfolded_tensors));
GraphDef folded_graph_def;
for (const std::pair<std::string, Tensor>& input : inputs) {
context.input_names.push_back(input.first);
}
context.output_names = outputs;
context.params["exclude_op"] = std::move(excluded_ops);
TF_ASSERT_OK(
graph_transforms::FoldConstants(graph_def, context, &folded_graph_def));
std::unique_ptr<tensorflow::Session> folded_session(
tensorflow::NewSession(tensorflow::SessionOptions()));
TF_ASSERT_OK(folded_session->Create(folded_graph_def));
std::vector<Tensor> folded_tensors;
TF_ASSERT_OK(folded_session->Run(inputs, outputs, {}, &folded_tensors));
EXPECT_EQ(unfolded_tensors.size(), folded_tensors.size());
for (int i = 0; i < unfolded_tensors.size(); ++i) {
test::ExpectTensorNear<float>(unfolded_tensors[i], folded_tensors[i],
1e-5);
}
std::map<std::string, const NodeDef*> folded_node_map;
for (const NodeDef& node : folded_graph_def.node()) {
folded_node_map.insert({node.name(), &node});
}
for (const NodeDef& node : graph_def.node()) {
const absl::string_view name(node.name());
const int occurrence_count = folded_node_map.count(node.name());
if (absl::EndsWith(name, "expect_removed")) {
EXPECT_EQ(0, occurrence_count) << "node.name()=" << node.name();
}
if (absl::EndsWith(name, "expect_remains")) {
EXPECT_EQ(1, occurrence_count) << "node.name()=" << node.name();
}
}
}
void TestReplaceSendRecvs() {
using namespace ::tensorflow::ops; // NOLINT(build/namespaces)
const int width = 100;
Tensor a_const_data(DT_FLOAT, TensorShape({width}));
test::FillIota<float>(&a_const_data, 1.0f);
auto o_root = tensorflow::Scope::NewRootScope();
_Recv(o_root.WithOpName("original_recv"), DT_FLOAT, "", "", 0, "");
Output o_a_const =
Const(o_root.WithOpName("a_const"), Input::Initializer(a_const_data));
Placeholder(o_root.WithOpName("placeholder"), DT_FLOAT);
_Send(o_root.WithOpName("original_send"), o_a_const, "", "", 0, "");
GraphDef o_graph_def;
TF_ASSERT_OK(o_root.ToGraphDef(&o_graph_def));
auto n_root = tensorflow::Scope::NewRootScope();
_Recv(n_root.WithOpName("original_recv"), DT_FLOAT, "", "", 0, "");
Output n_a_const =
Const(n_root.WithOpName("a_const"), Input::Initializer(a_const_data));
_Recv(n_root.WithOpName("_recv_placeholder_0"), DT_FLOAT, "", "", 0, "");
_Send(n_root.WithOpName("original_send"), n_a_const, "", "", 0, "");
_Send(n_root.WithOpName("new_send"), n_a_const, "", "", 0, "");
GraphDef n_graph_def;
TF_ASSERT_OK(n_root.ToGraphDef(&n_graph_def));
GraphDef result_graph_def;
TF_ASSERT_OK(graph_transforms::ReplaceSendRecvs(
o_graph_def, n_graph_def, {"placeholder"}, {"a_const"},
&result_graph_def));
std::map<std::string, const NodeDef*> node_map;
graph_transforms::MapNamesToNodes(result_graph_def, &node_map);
EXPECT_EQ(1, node_map.count("original_recv"));
EXPECT_EQ(1, node_map.count("a_const"));
EXPECT_EQ(1, node_map.count("placeholder"));
EXPECT_EQ(1, node_map.count("original_send"));
EXPECT_EQ(0, node_map.count("_recv_placeholder_0"));
EXPECT_EQ(0, node_map.count("new_send"));
}
void TestReplaceSendRecvsPrefixNames() {
using namespace ::tensorflow::ops; // NOLINT(build/namespaces)
auto o_root = tensorflow::Scope::NewRootScope();
auto a = Placeholder(o_root.WithOpName("placeholder"), DT_FLOAT);
auto b = Placeholder(o_root.WithOpName("placeholder_1"), DT_FLOAT);
auto add_o = Add(o_root.WithOpName("add"), a, b);
GraphDef o_graph_def;
TF_ASSERT_OK(o_root.ToGraphDef(&o_graph_def));
auto n_root = tensorflow::Scope::NewRootScope();
auto c = _Recv(n_root.WithOpName("_recv_placeholder_0"), DT_FLOAT, "", "",
0, "");
auto d = _Recv(n_root.WithOpName("_recv_placeholder_1_0"), DT_FLOAT, "", "",
0, "");
auto add_n = Add(n_root.WithOpName("add"), c, d);
GraphDef n_graph_def;
TF_ASSERT_OK(n_root.ToGraphDef(&n_graph_def));
GraphDef result_graph_def;
TF_ASSERT_OK(graph_transforms::ReplaceSendRecvs(
o_graph_def, n_graph_def, {"placeholder", "placeholder_1"}, {"add"},
&result_graph_def));
std::map<std::string, const NodeDef*> node_map;
graph_transforms::MapNamesToNodes(result_graph_def, &node_map);
EXPECT_EQ(1, node_map.count("placeholder"));
EXPECT_EQ(1, node_map.count("placeholder_1"));
EXPECT_EQ(1, node_map.count("add"));
}
void TestRemoveUnusedNodes() {
using namespace ::tensorflow::ops; // NOLINT(build/namespaces)
auto root = tensorflow::Scope::NewRootScope();
const int width = 100;
Tensor a_data(DT_FLOAT, TensorShape({width}));
test::FillIota<float>(&a_data, 1.0f);
Output a_const = Const(root.WithOpName("a"), Input::Initializer(a_data));
Tensor b_data(DT_FLOAT, TensorShape({width}));
test::FillIota<float>(&b_data, 1.0f);
Output b_const = Const(root.WithOpName("b"), Input::Initializer(b_data));
Output add = Add(root.WithOpName("add"), a_const, b_const);
Output placeholder = Placeholder(root.WithOpName("placeholder"), DT_FLOAT);
Output mul = Mul(root.WithOpName("output"), add, placeholder);
Tensor unused_data(DT_FLOAT, TensorShape({width}));
test::FillIota<float>(&unused_data, 1.0f);
Output unused_const =
Const(root.WithOpName("unused"), Input::Initializer(unused_data));
GraphDef graph_def;
TF_ASSERT_OK(root.ToGraphDef(&graph_def));
GraphDef result_graph_def;
TF_ASSERT_OK(graph_transforms::RemoveUnusedNodes(
graph_def, {{"placeholder"}, {"output"}}, &result_graph_def));
std::map<std::string, const NodeDef*> node_map;
graph_transforms::MapNamesToNodes(result_graph_def, &node_map);
EXPECT_EQ(1, node_map.count("a"));
EXPECT_EQ(1, node_map.count("b"));
EXPECT_EQ(1, node_map.count("add"));
EXPECT_EQ(1, node_map.count("placeholder"));
EXPECT_EQ(1, node_map.count("output"));
EXPECT_EQ(0, node_map.count("unused"));
}
void TestMaxConstantSizeInBytes() {
auto root = tensorflow::Scope::NewRootScope();
const int width = 100;
Tensor a_data(DT_FLOAT, TensorShape({width}));
test::FillIota<float>(&a_data, 1.0f);
Output a_const = ::tensorflow::ops::Const(
root.WithOpName("a_expect_remains"), Input::Initializer(a_data));
Tensor b_data(DT_FLOAT, TensorShape({width}));
test::FillIota<float>(&b_data, 1.0f);
Output b_const = ::tensorflow::ops::Const(
root.WithOpName("b_expect_remains"), Input::Initializer(b_data));
Output add = ::tensorflow::ops::Add(root.WithOpName("add_expect_remains"),
a_const, b_const);
Output placeholder = ::tensorflow::ops::Placeholder(
root.WithOpName("placeholder_expect_remains"), DT_FLOAT);
Output mul = ::tensorflow::ops::Mul(
root.WithOpName("output_expect_remains"), add, placeholder);
GraphDef graph_def;
TF_ASSERT_OK(root.ToGraphDef(&graph_def));
Tensor placeholder_tensor(DT_FLOAT, TensorShape({width}));
test::FillIota<float>(&placeholder_tensor, 1.0f);
// Setting the maximum constant size to 10 bytes should stop the constant
// folding at add(a, b) that would have yielded a constant of
// 100*sizeof(float) bytes.
graph_transforms::TransformFuncContext context;
context.params["max_constant_size_in_bytes"] = {"10"};
TestConstantFolding(graph_def,
{{"placeholder_expect_remains", placeholder_tensor}},
{}, {"output_expect_remains"}, context);
}
};
TEST_F(ConstantFoldingTest, TestSimpleAdd) { TestSimpleAdd(); }
TEST_F(ConstantFoldingTest, TestOpExclusionAdd) { TestOpExclusionAdd(); }
TEST_F(ConstantFoldingTest, TestShapePropagation) { TestShapePropagation(); }
TEST_F(ConstantFoldingTest, TestPreserveOutputShapes) {
TestPreserveOutputShapes();
}
TEST_F(ConstantFoldingTest, TestReplaceSendRecvs) { TestReplaceSendRecvs(); }
TEST_F(ConstantFoldingTest, TestReplaceSendRecvsPrefixNames) {
TestReplaceSendRecvsPrefixNames();
}
TEST_F(ConstantFoldingTest, TestRemoveUnusedNodes) { TestRemoveUnusedNodes(); }
TEST_F(ConstantFoldingTest, TestMaxConstantSizeInBytes) {
TestMaxConstantSizeInBytes();
}
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/fold_old_batch_norms.cc
+409
View File
@@ -0,0 +1,409 @@
/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/core/common_runtime/constant_folding.h"
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/graph/node_builder.h"
#include "tensorflow/core/graph/subgraph.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/tools/graph_transforms/fold_constants_lib.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
namespace {
// Ensures the tensor is the expected shape.
absl::Status ErrorIfNotVector(const Tensor& input,
const std::string& input_name,
int expected_width) {
if ((input.shape().dims() != 1) ||
(input.shape().dim_size(0) != expected_width)) {
return absl::InvalidArgumentError(absl::StrCat(
input_name,
" input to batch norm has bad shape: ", input.shape().DebugString()));
}
return absl::OkStatus();
}
absl::Status GetScaleAndOffsetValues(const NodeMatch& match,
std::vector<float>* scale_values,
std::vector<float>* offset_values) {
// Find all the nodes we expect in the subgraph.
const NodeDef& batch_norm_node = match.node;
// BatchNormWithGlobalNormalization and FusedBatchNorm ops only differ
// by input order and attribute names.
CHECK(batch_norm_node.op() == "BatchNormWithGlobalNormalization" ||
batch_norm_node.op() == "FusedBatchNorm");
const bool is_fused = batch_norm_node.op() == "FusedBatchNorm";
const int mean_idx = is_fused ? 3 : 1;
const int var_idx = is_fused ? 4 : 2;
const int beta_idx = is_fused ? 2 : 3;
const int gamma_idx = is_fused ? 1 : 4;
const std::string epsilon_attr = is_fused ? "epsilon" : "variance_epsilon";
// FusedBatchNorm always scales after normalization.
const bool scale_after_normalization =
is_fused || batch_norm_node.attr().at("scale_after_normalization").b();
const NodeDef& mean_node = match.inputs[mean_idx].node;
CHECK_EQ("Const", mean_node.op());
const NodeDef& variance_node = match.inputs[var_idx].node;
CHECK_EQ("Const", variance_node.op());
const NodeDef& beta_node = match.inputs[beta_idx].node;
CHECK_EQ("Const", beta_node.op());
const NodeDef& gamma_node = match.inputs[gamma_idx].node;
CHECK_EQ("Const", gamma_node.op());
// We have a set of vectors that we want to combine into a vector of
// scale values and offset values.
Tensor mean = GetNodeTensorAttr(mean_node, "value");
Tensor variance = GetNodeTensorAttr(variance_node, "value");
Tensor beta = GetNodeTensorAttr(beta_node, "value");
Tensor gamma = GetNodeTensorAttr(gamma_node, "value");
const float variance_epsilon = batch_norm_node.attr().at(epsilon_attr).f();
// Make sure all the inputs really are vectors with the same shape.
const int64_t num_cols = mean.shape().dim_size(0);
TF_RETURN_IF_ERROR(ErrorIfNotVector(variance, "Variance", num_cols));
TF_RETURN_IF_ERROR(ErrorIfNotVector(beta, "Beta", num_cols));
TF_RETURN_IF_ERROR(ErrorIfNotVector(gamma, "gamma", num_cols));
scale_values->resize(num_cols);
offset_values->resize(num_cols);
// Calculate the scale and offset values to apply.
if (scale_after_normalization) {
for (int i = 0; i < num_cols; ++i) {
(*scale_values)[i] =
(1.0f / sqrtf(variance.flat<float>()(i) + variance_epsilon)) *
gamma.flat<float>()(i);
}
} else {
for (int i = 0; i < num_cols; ++i) {
(*scale_values)[i] =
(1.0f / sqrtf(variance.flat<float>()(i) + variance_epsilon));
}
}
for (int i = 0; i < num_cols; ++i) {
(*offset_values)[i] =
(-mean.flat<float>()(i) * (*scale_values)[i]) + beta.flat<float>()(i);
}
return absl::OkStatus();
}
absl::Status FuseScaleOffsetToConvWeights(
const std::vector<float>& scale_values,
const std::vector<float>& offset_values, const NodeMatch& conv_node_match,
const std::string& conv_output_name, std::vector<NodeDef>* new_nodes) {
const NodeDef& conv_node = conv_node_match.node;
// CHECK_EQ("Conv2D", conv_node.op());
const NodeDef& input_node = conv_node_match.inputs[0].node;
const NodeDef& weights_node = conv_node_match.inputs[1].node;
CHECK_EQ("Const", weights_node.op());
Tensor weights = GetNodeTensorAttr(weights_node, "value");
int64_t weights_cols;
if (conv_node.op() == "Conv2D") {
weights_cols = weights.shape().dim_size(3);
} else if (conv_node.op() == "DepthwiseConv2dNative") {
weights_cols = weights.shape().dim_size(2) * weights.shape().dim_size(3);
} else {
weights_cols = weights.shape().dim_size(1);
}
CHECK_EQ(weights_cols, scale_values.size());
// Multiply the original weights by the scale vector.
auto weights_vector = weights.flat<float>();
Tensor scaled_weights(DT_FLOAT, weights.shape());
auto scaled_weights_vector = scaled_weights.flat<float>();
for (int64_t row = 0; row < weights_vector.dimension(0); ++row) {
scaled_weights_vector(row) =
weights_vector(row) * scale_values[row % weights_cols];
}
// Figure out the remaining bias to add on.
Tensor bias_offset(DT_FLOAT, {weights_cols});
auto bias_offset_vector = bias_offset.flat<float>();
for (int64_t col = 0; col < weights_cols; ++col) {
bias_offset_vector(col) = offset_values[col];
}
// Construct the new nodes.
NodeDef scaled_weights_node;
scaled_weights_node.set_op("Const");
scaled_weights_node.set_name(weights_node.name());
SetNodeAttr("dtype", DT_FLOAT, &scaled_weights_node);
SetNodeTensorAttr<float>("value", scaled_weights, &scaled_weights_node);
new_nodes->push_back(scaled_weights_node);
// The input and convolution can be copied straight over, since the
// name of the scaled weights constant is the same as the original.
new_nodes->push_back(input_node);
new_nodes->push_back(conv_node);
NodeDef bias_offset_node;
bias_offset_node.set_op("Const");
bias_offset_node.set_name(conv_node.name() + "_bn_offset");
SetNodeAttr("dtype", DT_FLOAT, &bias_offset_node);
SetNodeTensorAttr<float>("value", bias_offset, &bias_offset_node);
new_nodes->push_back(bias_offset_node);
NodeDef bias_add_node;
bias_add_node.set_op("BiasAdd");
bias_add_node.set_name(conv_output_name);
if (conv_node.attr().count("data_format")) {
CopyNodeAttr(conv_node, "data_format", "data_format", &bias_add_node);
}
CopyNodeAttr(conv_node, "T", "T", &bias_add_node);
AddNodeInput(conv_node.name(), &bias_add_node);
AddNodeInput(bias_offset_node.name(), &bias_add_node);
new_nodes->push_back(bias_add_node);
return absl::OkStatus();
}
absl::Status FuseBatchNormWithConv(const NodeMatch& match,
std::vector<NodeDef>* new_nodes) {
// Calculate the scale and offset values to apply.
std::vector<float> scale_values;
std::vector<float> offset_values;
TF_RETURN_IF_ERROR(
GetScaleAndOffsetValues(match, &scale_values, &offset_values));
// Fuse conv weights, and set the final output node name as batch_norm_node.
const NodeDef& batch_norm_node = match.node;
TF_RETURN_IF_ERROR(
FuseScaleOffsetToConvWeights(scale_values, offset_values, match.inputs[0],
batch_norm_node.name(), new_nodes));
return absl::OkStatus();
}
absl::Status FuseBatchNormWithBatchToSpace(const NodeMatch& match,
std::vector<NodeDef>* new_nodes) {
// Calculate the scale and offset values to apply.
std::vector<float> scale_values;
std::vector<float> offset_values;
TF_RETURN_IF_ERROR(
GetScaleAndOffsetValues(match, &scale_values, &offset_values));
// Fuse conv weights, and set the final output node name as batch_norm_node.
const NodeDef& batch_norm_node = match.node;
const NodeMatch& batch_to_space_node_match = match.inputs[0];
const NodeMatch& conv_node_match = batch_to_space_node_match.inputs[0];
const NodeDef& batch_to_space_node = batch_to_space_node_match.node;
const NodeDef& conv_node = conv_node_match.node;
std::string biasadd_name = conv_node.name() + "/biasadd";
TF_RETURN_IF_ERROR(FuseScaleOffsetToConvWeights(
scale_values, offset_values, conv_node_match, biasadd_name, new_nodes));
NodeDef new_batch_to_space_node = batch_to_space_node;
// reuse batch_norm node name
new_batch_to_space_node.set_name(batch_norm_node.name());
new_batch_to_space_node.set_input(0, biasadd_name);
new_nodes->push_back(batch_to_space_node_match.inputs[1].node);
new_nodes->push_back(batch_to_space_node_match.inputs[2].node);
new_nodes->push_back(new_batch_to_space_node);
return absl::OkStatus();
}
absl::Status FuseBatchNormWithConvConcat(const NodeMatch& match,
std::vector<NodeDef>* new_nodes) {
// Calculate the scale and offset values to apply.
std::vector<float> scale_values;
std::vector<float> offset_values;
TF_RETURN_IF_ERROR(
GetScaleAndOffsetValues(match, &scale_values, &offset_values));
// Find all the nodes we expect in the subgraph.
const NodeDef& batch_norm_node = match.node;
const NodeMatch& concat_node_match = match.inputs[0];
NodeDef concat_node = concat_node_match.node;
CHECK_EQ("ConcatV2", concat_node.op());
// First process the axis.
NodeDef axis_node = concat_node_match.inputs[2].node;
CHECK_EQ("Const", axis_node.op());
Tensor axis = GetNodeTensorAttr(axis_node, "value");
int32_t axis_scalar = (axis.scalar<int32_t>())();
// Set both conv0 and conv1 have the same scale and offset in default.
std::vector<float> scale0(scale_values);
std::vector<float> offset0(offset_values);
std::vector<float> scale1(scale_values);
std::vector<float> offset1(offset_values);
if (axis_scalar == 3) {
// If axis is 3, then scale and offset will be split into two halfs.
const NodeDef& weights0_node = concat_node_match.inputs[0].inputs[1].node;
Tensor weights0 = GetNodeTensorAttr(weights0_node, "value");
const int64_t split_cols = weights0.shape().dim_size(3);
// Only keep the first half for scale0/offset0.
scale0.erase(scale0.begin() + split_cols, scale0.end());
offset0.erase(offset0.begin() + split_cols, offset0.end());
// Only keep the second half for scale1/offset1.
scale1.erase(scale1.begin(), scale1.begin() + split_cols);
offset1.erase(offset1.begin(), offset1.begin() + split_cols);
}
// Fuse the weights for input0 of conv2d.
const std::string concat0_output_name = concat_node.name() + "_bn_in0";
TF_RETURN_IF_ERROR(
FuseScaleOffsetToConvWeights(scale0, offset0, concat_node_match.inputs[0],
concat0_output_name, new_nodes));
// Fuse the weights for input1 of conv2d.
const std::string concat1_output_name = concat_node.name() + "_bn_in1";
TF_RETURN_IF_ERROR(
FuseScaleOffsetToConvWeights(scale1, offset1, concat_node_match.inputs[1],
concat1_output_name, new_nodes));
// Push the shape node.
new_nodes->push_back(concat_node_match.inputs[2].node);
// Set the final output op name to batch_normal_node.
concat_node.set_name(batch_norm_node.name());
concat_node.set_input(0, concat0_output_name);
concat_node.set_input(1, concat1_output_name);
new_nodes->push_back(concat_node);
return absl::OkStatus();
}
} // namespace
// Finds monolithic batch norm ops (as used in early versions of TensorFlow) and
// converts them into premultiplied weight inputs to convolutions.
absl::Status FoldOldBatchNorms(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def) {
GraphDef current_graph_def = input_graph_def;
// We have to do several passes to catch all the old BN nodes, since many of
// them may share inputs and so be excluded from replacement in one pass.
bool did_graph_change;
do {
did_graph_change = false;
GraphDef replaced_graph_def;
TF_RETURN_IF_ERROR(ReplaceMatchingOpTypes(
current_graph_def, // clang-format off
{"BatchNormWithGlobalNormalization|FusedBatchNorm", // batch_norm_node
{
{"Conv2D|DepthwiseConv2dNative", // conv_node
{
{"*"}, // input_node
{"Const"}, // weights_node
}
},
{"Const"}, // mean_node
{"Const"}, // variance_node
{"Const"}, // beta_node
{"Const"}, // gamma_node
}
}, // clang-format on
[&did_graph_change](const NodeMatch& match,
const std::set<std::string>& input_nodes,
const std::set<std::string>& output_nodes,
std::vector<NodeDef>* new_nodes) {
TF_RETURN_IF_ERROR(FuseBatchNormWithConv(match, new_nodes));
did_graph_change = true;
return absl::OkStatus();
},
{}, &replaced_graph_def));
current_graph_def = replaced_graph_def;
} while (did_graph_change);
do {
did_graph_change = false;
GraphDef replaced_graph_def;
TF_RETURN_IF_ERROR(ReplaceMatchingOpTypes(
current_graph_def, // clang-format off
{"BatchNormWithGlobalNormalization|FusedBatchNorm", // batch_norm_node
{
{"BatchToSpaceND", // batch_to_space_node
{
{"Conv2D|DepthwiseConv2dNative", // conv_node
{
{"*"}, // input_node
{"Const"}, // weights_node
}
},
{"Const"}, // block_shape
{"Const"}, // crops
}
},
{"Const"}, // mean_node
{"Const"}, // variance_node
{"Const"}, // beta_node
{"Const"}, // gamma_node
}
}, // clang-format on
[&did_graph_change](const NodeMatch& match,
const std::set<std::string>& input_nodes,
const std::set<std::string>& output_nodes,
std::vector<NodeDef>* new_nodes) {
TF_RETURN_IF_ERROR(FuseBatchNormWithBatchToSpace(match, new_nodes));
did_graph_change = true;
return absl::OkStatus();
},
{}, &replaced_graph_def));
current_graph_def = replaced_graph_def;
} while (did_graph_change);
do {
did_graph_change = false;
GraphDef replaced_graph_def;
// Replace BatchNorm with concat as input.
TF_RETURN_IF_ERROR(ReplaceMatchingOpTypes(
current_graph_def, // clang-format off
{"BatchNormWithGlobalNormalization|FusedBatchNorm", // batch_norm_node
{
{"ConcatV2|Concat", // concat two conv2d.
{
{"Conv2D|DepthwiseConv2dNative", // conv_node
{
{"*"}, // input_node
{"Const"}, // weights_node
}
},
{"Conv2D|DepthwiseConv2dNative", // conv_node
{
{"*"}, // input_node
{"Const"}, // weights_node
}
},
{"Const"}, // axis
},
},
{"Const"}, // mean_node
{"Const"}, // variance_node
{"Const"}, // beta_node
{"Const"}, // gamma_node
}
}, // clang-format on
[&did_graph_change](const NodeMatch& match,
const std::set<std::string>& input_nodes,
const std::set<std::string>& output_nodes,
std::vector<NodeDef>* new_nodes) {
TF_RETURN_IF_ERROR(FuseBatchNormWithConvConcat(match, new_nodes));
did_graph_change = true;
return absl::OkStatus();
},
{}, &replaced_graph_def));
current_graph_def = replaced_graph_def;
} while (did_graph_change);
*output_graph_def = current_graph_def;
return absl::OkStatus();
}
REGISTER_GRAPH_TRANSFORM("fold_old_batch_norms", FoldOldBatchNorms);
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/fold_old_batch_norms_test.cc
+580
View File
@@ -0,0 +1,580 @@
/* Copyright 2015 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/cc/ops/const_op.h"
#include "tensorflow/cc/ops/image_ops.h"
#include "tensorflow/cc/ops/nn_ops.h"
#include "tensorflow/cc/ops/array_ops.h"
#include "tensorflow/cc/ops/sendrecv_ops.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/framework/versions.pb.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
// Declare here, so we don't need a public header.
absl::Status FoldOldBatchNorms(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def);
class FoldOldBatchNormsTest : public ::testing::Test {
protected:
void TestFoldOldBatchNorms() {
auto root = tensorflow::Scope::NewRootScope();
using namespace ::tensorflow::ops; // NOLINT(build/namespaces)
Tensor input_data(DT_FLOAT, TensorShape({1, 1, 6, 2}));
test::FillValues<float>(
&input_data, {1.0f, 4.0f, 2.0f, 5.0f, 3.0f, 6.0f, -1.0f, -4.0f, -2.0f,
-5.0f, -3.0f, -6.0f});
Output input_op =
Const(root.WithOpName("input_op"), Input::Initializer(input_data));
Tensor weights_data(DT_FLOAT, TensorShape({1, 2, 2, 2}));
test::FillValues<float>(&weights_data,
{1.0f, 2.0f, 3.0f, 4.0f, 0.1f, 0.2f, 0.3f, 0.4f});
Output weights_op =
Const(root.WithOpName("weights_op"), Input::Initializer(weights_data));
Output conv_op = Conv2D(root.WithOpName("conv_op"), input_op, weights_op,
{1, 1, 1, 1}, "VALID");
Tensor mean_data(DT_FLOAT, TensorShape({2}));
test::FillValues<float>(&mean_data, {10.0f, 20.0f});
Output mean_op =
Const(root.WithOpName("mean_op"), Input::Initializer(mean_data));
Tensor variance_data(DT_FLOAT, TensorShape({2}));
test::FillValues<float>(&variance_data, {0.25f, 0.5f});
Output variance_op = Const(root.WithOpName("variance_op"),
Input::Initializer(variance_data));
Tensor beta_data(DT_FLOAT, TensorShape({2}));
test::FillValues<float>(&beta_data, {0.1f, 0.6f});
Output beta_op =
Const(root.WithOpName("beta_op"), Input::Initializer(beta_data));
Tensor gamma_data(DT_FLOAT, TensorShape({2}));
test::FillValues<float>(&gamma_data, {1.0f, 2.0f});
Output gamma_op =
Const(root.WithOpName("gamma_op"), Input::Initializer(gamma_data));
GraphDef original_graph_def;
TF_ASSERT_OK(root.ToGraphDef(&original_graph_def));
// This is needed because we're trying to convert over a deprecated op which
// should only be present in older GraphDef files. Without this we see a
// deprecation error.
// This is justified because we're trying to test a tool that is expected to
// run on legacy files, to help users convert over to less problematic
// versions.
NodeDef batch_norm_node;
batch_norm_node.set_op("BatchNormWithGlobalNormalization");
batch_norm_node.set_name("output");
AddNodeInput("conv_op", &batch_norm_node);
AddNodeInput("mean_op", &batch_norm_node);
AddNodeInput("variance_op", &batch_norm_node);
AddNodeInput("beta_op", &batch_norm_node);
AddNodeInput("gamma_op", &batch_norm_node);
SetNodeAttr("T", DT_FLOAT, &batch_norm_node);
SetNodeAttr("variance_epsilon", 0.00001f, &batch_norm_node);
SetNodeAttr("scale_after_normalization", false, &batch_norm_node);
*(original_graph_def.mutable_node()->Add()) = batch_norm_node;
original_graph_def.mutable_versions()->set_producer(8);
std::unique_ptr<Session> original_session(NewSession(SessionOptions()));
TF_ASSERT_OK(original_session->Create(original_graph_def));
std::vector<Tensor> original_outputs;
TF_ASSERT_OK(original_session->Run({}, {"output"}, {}, &original_outputs));
GraphDef fused_graph_def;
TF_ASSERT_OK(FoldOldBatchNorms(original_graph_def, {{}, {"output"}},
&fused_graph_def));
std::unique_ptr<Session> fused_session(NewSession(SessionOptions()));
TF_ASSERT_OK(fused_session->Create(fused_graph_def));
std::vector<Tensor> fused_outputs;
TF_ASSERT_OK(fused_session->Run({}, {"output"}, {}, &fused_outputs));
test::ExpectTensorNear<float>(original_outputs[0], fused_outputs[0], 1e-5);
for (const NodeDef& node : fused_graph_def.node()) {
EXPECT_NE("BatchNormWithGlobalNormalization", node.op());
}
}
void TestFoldOldBatchNormsAfterDepthwiseConv2dNative() {
auto root = tensorflow::Scope::NewRootScope();
using namespace ::tensorflow::ops; // NOLINT(build/namespaces)
Tensor input_data(DT_FLOAT, TensorShape({1, 1, 6, 2}));
test::FillValues<float>(
&input_data, {1.0f, 4.0f, 2.0f, 5.0f, 3.0f, 6.0f, -1.0f, -4.0f, -2.0f,
-5.0f, -3.0f, -6.0f});
Output input_op =
Const(root.WithOpName("input_op"), Input::Initializer(input_data));
Tensor weights_data(DT_FLOAT, TensorShape({1, 2, 2, 2}));
test::FillValues<float>(&weights_data,
{1.0f, 2.0f, 3.0f, 4.0f, 0.1f, 0.2f, 0.3f, 0.4f});
Output weights_op =
Const(root.WithOpName("weights_op"), Input::Initializer(weights_data));
Output conv_op = DepthwiseConv2dNative(root.WithOpName("conv_op"), input_op,
weights_op, {1, 1, 1, 1}, "VALID");
Tensor mean_data(DT_FLOAT, TensorShape({4}));
test::FillValues<float>(&mean_data, {10.0f, 20.0f, 30.0f, 40.0f});
Output mean_op =
Const(root.WithOpName("mean_op"), Input::Initializer(mean_data));
Tensor variance_data(DT_FLOAT, TensorShape({4}));
test::FillValues<float>(&variance_data, {0.25f, 0.5f, 0.75f, 1.0f});
Output variance_op = Const(root.WithOpName("variance_op"),
Input::Initializer(variance_data));
Tensor beta_data(DT_FLOAT, TensorShape({4}));
test::FillValues<float>(&beta_data, {0.1f, 0.6f, 1.1f, 1.6f});
Output beta_op =
Const(root.WithOpName("beta_op"), Input::Initializer(beta_data));
Tensor gamma_data(DT_FLOAT, TensorShape({4}));
test::FillValues<float>(&gamma_data, {1.0f, 2.0f, 3.0f, 4.0f});
Output gamma_op =
Const(root.WithOpName("gamma_op"), Input::Initializer(gamma_data));
GraphDef original_graph_def;
TF_ASSERT_OK(root.ToGraphDef(&original_graph_def));
NodeDef batch_norm_node;
batch_norm_node.set_op("BatchNormWithGlobalNormalization");
batch_norm_node.set_name("output");
AddNodeInput("conv_op", &batch_norm_node);
AddNodeInput("mean_op", &batch_norm_node);
AddNodeInput("variance_op", &batch_norm_node);
AddNodeInput("beta_op", &batch_norm_node);
AddNodeInput("gamma_op", &batch_norm_node);
SetNodeAttr("T", DT_FLOAT, &batch_norm_node);
SetNodeAttr("variance_epsilon", 0.00001f, &batch_norm_node);
SetNodeAttr("scale_after_normalization", false, &batch_norm_node);
*(original_graph_def.mutable_node()->Add()) = batch_norm_node;
original_graph_def.mutable_versions()->set_producer(8);
std::unique_ptr<Session> original_session(NewSession(SessionOptions()));
TF_ASSERT_OK(original_session->Create(original_graph_def));
std::vector<Tensor> original_outputs;
TF_ASSERT_OK(original_session->Run({}, {"output"}, {}, &original_outputs));
GraphDef fused_graph_def;
TF_ASSERT_OK(FoldOldBatchNorms(original_graph_def, {{}, {"output"}},
&fused_graph_def));
std::unique_ptr<Session> fused_session(NewSession(SessionOptions()));
TF_ASSERT_OK(fused_session->Create(fused_graph_def));
std::vector<Tensor> fused_outputs;
TF_ASSERT_OK(fused_session->Run({}, {"output"}, {}, &fused_outputs));
test::ExpectTensorNear<float>(original_outputs[0], fused_outputs[0], 1e-5);
for (const NodeDef& node : fused_graph_def.node()) {
EXPECT_NE("BatchNormWithGlobalNormalization", node.op());
}
}
void TestFoldFusedBatchNorms() {
auto root = tensorflow::Scope::NewRootScope();
using namespace ::tensorflow::ops; // NOLINT(build/namespaces)
Tensor input_data(DT_FLOAT, TensorShape({1, 1, 6, 2}));
test::FillValues<float>(
&input_data, {1.0f, 4.0f, 2.0f, 5.0f, 3.0f, 6.0f, -1.0f, -4.0f, -2.0f,
-5.0f, -3.0f, -6.0f});
Output input_op =
Const(root.WithOpName("input_op"), Input::Initializer(input_data));
Tensor weights_data(DT_FLOAT, TensorShape({1, 2, 2, 2}));
test::FillValues<float>(&weights_data,
{1.0f, 2.0f, 3.0f, 4.0f, 0.1f, 0.2f, 0.3f, 0.4f});
Output weights_op =
Const(root.WithOpName("weights_op"), Input::Initializer(weights_data));
Output conv_op = Conv2D(root.WithOpName("conv_op"), input_op, weights_op,
{1, 1, 1, 1}, "VALID");
Tensor mean_data(DT_FLOAT, TensorShape({2}));
test::FillValues<float>(&mean_data, {10.0f, 20.0f});
Output mean_op =
Const(root.WithOpName("mean_op"), Input::Initializer(mean_data));
Tensor variance_data(DT_FLOAT, TensorShape({2}));
test::FillValues<float>(&variance_data, {0.25f, 0.5f});
Output variance_op = Const(root.WithOpName("variance_op"),
Input::Initializer(variance_data));
Tensor beta_data(DT_FLOAT, TensorShape({2}));
test::FillValues<float>(&beta_data, {0.1f, 0.6f});
Output beta_op =
Const(root.WithOpName("beta_op"), Input::Initializer(beta_data));
Tensor gamma_data(DT_FLOAT, TensorShape({2}));
test::FillValues<float>(&gamma_data, {1.0f, 2.0f});
Output gamma_op =
Const(root.WithOpName("gamma_op"), Input::Initializer(gamma_data));
GraphDef original_graph_def;
TF_ASSERT_OK(root.ToGraphDef(&original_graph_def));
NodeDef batch_norm_node;
batch_norm_node.set_op("FusedBatchNorm");
batch_norm_node.set_name("output");
AddNodeInput("conv_op", &batch_norm_node);
AddNodeInput("gamma_op", &batch_norm_node);
AddNodeInput("beta_op", &batch_norm_node);
AddNodeInput("mean_op", &batch_norm_node);
AddNodeInput("variance_op", &batch_norm_node);
SetNodeAttr("T", DT_FLOAT, &batch_norm_node);
SetNodeAttr("epsilon", 0.00001f, &batch_norm_node);
SetNodeAttr("is_training", false, &batch_norm_node);
*(original_graph_def.mutable_node()->Add()) = batch_norm_node;
std::unique_ptr<Session> original_session(NewSession(SessionOptions()));
TF_ASSERT_OK(original_session->Create(original_graph_def));
std::vector<Tensor> original_outputs;
TF_ASSERT_OK(original_session->Run({}, {"output"}, {}, &original_outputs));
GraphDef fused_graph_def;
TF_ASSERT_OK(FoldOldBatchNorms(original_graph_def, {{}, {"output"}},
&fused_graph_def));
std::unique_ptr<Session> fused_session(NewSession(SessionOptions()));
TF_ASSERT_OK(fused_session->Create(fused_graph_def));
std::vector<Tensor> fused_outputs;
TF_ASSERT_OK(fused_session->Run({}, {"output"}, {}, &fused_outputs));
test::ExpectTensorNear<float>(original_outputs[0], fused_outputs[0], 2e-5);
for (const NodeDef& node : fused_graph_def.node()) {
EXPECT_NE("FusedBatchNorm", node.op());
}
}
void TestFoldFusedBatchNormsAfterDepthwiseConv2dNative() {
auto root = tensorflow::Scope::NewRootScope();
using namespace ::tensorflow::ops; // NOLINT(build/namespaces)
Tensor input_data(DT_FLOAT, TensorShape({1, 1, 6, 2}));
test::FillValues<float>(
&input_data, {1.0f, 4.0f, 2.0f, 5.0f, 3.0f, 6.0f, -1.0f, -4.0f, -2.0f,
-5.0f, -3.0f, -6.0f});
Output input_op =
Const(root.WithOpName("input_op"), Input::Initializer(input_data));
Tensor weights_data(DT_FLOAT, TensorShape({1, 2, 2, 2}));
test::FillValues<float>(&weights_data,
{1.0f, 2.0f, 3.0f, 4.0f, 0.1f, 0.2f, 0.3f, 0.4f});
Output weights_op =
Const(root.WithOpName("weights_op"), Input::Initializer(weights_data));
Output conv_op = DepthwiseConv2dNative(root.WithOpName("conv_op"), input_op,
weights_op, {1, 1, 1, 1}, "VALID");
Tensor mean_data(DT_FLOAT, TensorShape({4}));
test::FillValues<float>(&mean_data, {10.0f, 20.0f, 30.0f, 40.0f});
Output mean_op =
Const(root.WithOpName("mean_op"), Input::Initializer(mean_data));
Tensor variance_data(DT_FLOAT, TensorShape({4}));
test::FillValues<float>(&variance_data, {0.25f, 0.5f, 0.75f, 1.0f});
Output variance_op = Const(root.WithOpName("variance_op"),
Input::Initializer(variance_data));
Tensor beta_data(DT_FLOAT, TensorShape({4}));
test::FillValues<float>(&beta_data, {0.1f, 0.6f, 1.1f, 1.6f});
Output beta_op =
Const(root.WithOpName("beta_op"), Input::Initializer(beta_data));
Tensor gamma_data(DT_FLOAT, TensorShape({4}));
test::FillValues<float>(&gamma_data, {1.0f, 2.0f, 3.0f, 4.0f});
Output gamma_op =
Const(root.WithOpName("gamma_op"), Input::Initializer(gamma_data));
GraphDef original_graph_def;
TF_ASSERT_OK(root.ToGraphDef(&original_graph_def));
NodeDef batch_norm_node;
batch_norm_node.set_op("FusedBatchNorm");
batch_norm_node.set_name("output");
AddNodeInput("conv_op", &batch_norm_node);
AddNodeInput("gamma_op", &batch_norm_node);
AddNodeInput("beta_op", &batch_norm_node);
AddNodeInput("mean_op", &batch_norm_node);
AddNodeInput("variance_op", &batch_norm_node);
SetNodeAttr("T", DT_FLOAT, &batch_norm_node);
SetNodeAttr("epsilon", 0.00001f, &batch_norm_node);
SetNodeAttr("is_training", false, &batch_norm_node);
*(original_graph_def.mutable_node()->Add()) = batch_norm_node;
std::unique_ptr<Session> original_session(NewSession(SessionOptions()));
TF_ASSERT_OK(original_session->Create(original_graph_def));
std::vector<Tensor> original_outputs;
TF_ASSERT_OK(original_session->Run({}, {"output"}, {}, &original_outputs));
GraphDef fused_graph_def;
TF_ASSERT_OK(FoldOldBatchNorms(original_graph_def, {{}, {"output"}},
&fused_graph_def));
std::unique_ptr<Session> fused_session(NewSession(SessionOptions()));
TF_ASSERT_OK(fused_session->Create(fused_graph_def));
std::vector<Tensor> fused_outputs;
TF_ASSERT_OK(fused_session->Run({}, {"output"}, {}, &fused_outputs));
test::ExpectClose(original_outputs[0], fused_outputs[0], /*atol=*/2e-5,
/*rtol=*/2e-5);
for (const NodeDef& node : fused_graph_def.node()) {
EXPECT_NE("FusedBatchNorm", node.op());
}
}
void TestFoldFusedBatchNormsWithConcat(const bool split) {
auto root = tensorflow::Scope::NewRootScope();
using namespace ::tensorflow::ops; // NOLINT(build/namespaces)
// If split is true, concat two inputs at dim=3; otherwise, concat at dim 2.
auto input_shape =
split ? TensorShape({1, 1, 6, 2}) : TensorShape({1, 1, 12, 1});
Tensor input_data(DT_FLOAT, input_shape);
test::FillValues<float>(
&input_data, {1.0f, 4.0f, 2.0f, 5.0f, 3.0f, 6.0f, -1.0f, -4.0f, -2.0f,
-5.0f, -3.0f, -6.0f});
Output input0_op =
Const(root.WithOpName("input_op0"), Input::Initializer(input_data));
// If split is true, concat two inputs at dim=3; otherwise, concat at dim 2.
// The final output shape of concat is always {1, 2, 2, 2}.
auto weight_shape =
split ? TensorShape({1, 2, 2, 1}) : TensorShape({1, 2, 1, 2});
Tensor weights0_data(DT_FLOAT, weight_shape);
test::FillValues<float>(&weights0_data, {1.0f, 2.0f, 3.0f, 4.0f});
Output weights0_op = Const(root.WithOpName("weights1_op"),
Input::Initializer(weights0_data));
Output conv0_op = Conv2D(root.WithOpName("conv1_op"), input0_op,
weights0_op, {1, 1, 1, 1}, "VALID");
Output input1_op =
Const(root.WithOpName("input1_op"), Input::Initializer(input_data));
Tensor weights1_data(DT_FLOAT, weight_shape);
test::FillValues<float>(&weights1_data, {1.0f, 2.0f, 3.0f, 4.0f});
Output weights1_op = Const(root.WithOpName("weights1_op"),
Input::Initializer(weights1_data));
Output conv1_op = Conv2D(root.WithOpName("conv1_op"), input1_op,
weights1_op, {1, 1, 1, 1}, "VALID");
Tensor shape_tensor(DT_INT32, TensorShape({}));
// Concat at dim 3 if split; otherwise, concat at dim 2.
int32_t concat_axis = split ? 3 : 2;
test::FillValues<int32_t>(&shape_tensor, {concat_axis});
Output shape_op =
Const(root.WithOpName("shape_op"), Input::Initializer(shape_tensor));
Output concat_op =
Concat(root.WithOpName("concat_op"), {conv0_op, conv1_op}, shape_op);
Tensor mean_data(DT_FLOAT, TensorShape({2}));
test::FillValues<float>(&mean_data, {10.0f, 20.0f});
Output mean_op =
Const(root.WithOpName("mean_op"), Input::Initializer(mean_data));
Tensor variance_data(DT_FLOAT, TensorShape({2}));
test::FillValues<float>(&variance_data, {0.25f, 0.5f});
Output variance_op = Const(root.WithOpName("variance_op"),
Input::Initializer(variance_data));
Tensor beta_data(DT_FLOAT, TensorShape({2}));
test::FillValues<float>(&beta_data, {0.1f, 0.6f});
Output beta_op =
Const(root.WithOpName("beta_op"), Input::Initializer(beta_data));
Tensor gamma_data(DT_FLOAT, TensorShape({2}));
test::FillValues<float>(&gamma_data, {1.0f, 2.0f});
Output gamma_op =
Const(root.WithOpName("gamma_op"), Input::Initializer(gamma_data));
GraphDef original_graph_def;
TF_ASSERT_OK(root.ToGraphDef(&original_graph_def));
NodeDef batch_norm_node;
batch_norm_node.set_op("FusedBatchNorm");
batch_norm_node.set_name("output");
AddNodeInput("concat_op", &batch_norm_node);
AddNodeInput("gamma_op", &batch_norm_node);
AddNodeInput("beta_op", &batch_norm_node);
AddNodeInput("mean_op", &batch_norm_node);
AddNodeInput("variance_op", &batch_norm_node);
SetNodeAttr("T", DT_FLOAT, &batch_norm_node);
SetNodeAttr("epsilon", 0.00001f, &batch_norm_node);
SetNodeAttr("is_training", false, &batch_norm_node);
*(original_graph_def.mutable_node()->Add()) = batch_norm_node;
std::unique_ptr<Session> original_session(NewSession(SessionOptions()));
TF_ASSERT_OK(original_session->Create(original_graph_def));
std::vector<Tensor> original_outputs;
TF_ASSERT_OK(original_session->Run({}, {"output"}, {}, &original_outputs));
GraphDef fused_graph_def;
TF_ASSERT_OK(FoldOldBatchNorms(original_graph_def, {{}, {"output"}},
&fused_graph_def));
std::unique_ptr<Session> fused_session(NewSession(SessionOptions()));
TF_ASSERT_OK(fused_session->Create(fused_graph_def));
std::vector<Tensor> fused_outputs;
TF_ASSERT_OK(fused_session->Run({}, {"output"}, {}, &fused_outputs));
test::ExpectClose(original_outputs[0], fused_outputs[0]);
for (const NodeDef& node : fused_graph_def.node()) {
EXPECT_NE("FusedBatchNorm", node.op());
}
}
};
void TestFoldFusedBatchNormsWithBatchToSpace() {
auto root = tensorflow::Scope::NewRootScope();
using namespace ::tensorflow::ops; // NOLINT(build/namespaces)
Tensor input_data(DT_FLOAT, TensorShape({2, 1, 3, 2}));
test::FillValues<float>(
&input_data, {1.0f, 4.0f, 2.0f, 5.0f, 3.0f, 6.0f, -1.0f, -4.0f, -2.0f,
-5.0f, -3.0f, -6.0f});
Output input_op =
Const(root.WithOpName("input_op"), Input::Initializer(input_data));
Tensor weights_data(DT_FLOAT, TensorShape({1, 2, 2, 2}));
test::FillValues<float>(&weights_data,
{1.0f, 2.0f, 3.0f, 4.0f, 0.1f, 0.2f, 0.3f, 0.4f});
Output weights_op =
Const(root.WithOpName("weights_op"), Input::Initializer(weights_data));
Output conv_op = Conv2D(root.WithOpName("conv_op"), input_op, weights_op,
{1, 1, 1, 1}, "VALID");
Tensor block_shape_data(DT_INT32, TensorShape({2}));
test::FillValues<int32_t>(&block_shape_data, {1, 2});
Output block_shape_op = Const(root.WithOpName("block_shape_op"),
Input::Initializer(block_shape_data));
Tensor crops_data(DT_INT32, TensorShape({2, 2}));
test::FillValues<int32_t>(&crops_data, {0, 0, 0, 1});
Output crops_op =
Const(root.WithOpName("crops_op"), Input::Initializer(crops_data));
Output batch_to_space_op =
BatchToSpaceND(root.WithOpName("batch_to_space_op"), conv_op,
block_shape_op, crops_data);
Tensor mean_data(DT_FLOAT, TensorShape({2}));
test::FillValues<float>(&mean_data, {10.0f, 20.0f});
Output mean_op =
Const(root.WithOpName("mean_op"), Input::Initializer(mean_data));
Tensor variance_data(DT_FLOAT, TensorShape({2}));
test::FillValues<float>(&variance_data, {0.25f, 0.5f});
Output variance_op =
Const(root.WithOpName("variance_op"), Input::Initializer(variance_data));
Tensor beta_data(DT_FLOAT, TensorShape({2}));
test::FillValues<float>(&beta_data, {0.1f, 0.6f});
Output beta_op =
Const(root.WithOpName("beta_op"), Input::Initializer(beta_data));
Tensor gamma_data(DT_FLOAT, TensorShape({2}));
test::FillValues<float>(&gamma_data, {1.0f, 2.0f});
Output gamma_op =
Const(root.WithOpName("gamma_op"), Input::Initializer(gamma_data));
GraphDef original_graph_def;
TF_ASSERT_OK(root.ToGraphDef(&original_graph_def));
NodeDef batch_norm_node;
batch_norm_node.set_op("FusedBatchNorm");
batch_norm_node.set_name("output");
AddNodeInput("batch_to_space_op", &batch_norm_node);
AddNodeInput("gamma_op", &batch_norm_node);
AddNodeInput("beta_op", &batch_norm_node);
AddNodeInput("mean_op", &batch_norm_node);
AddNodeInput("variance_op", &batch_norm_node);
SetNodeAttr("T", DT_FLOAT, &batch_norm_node);
SetNodeAttr("epsilon", 0.00001f, &batch_norm_node);
SetNodeAttr("is_training", false, &batch_norm_node);
*(original_graph_def.mutable_node()->Add()) = batch_norm_node;
std::unique_ptr<Session> original_session(NewSession(SessionOptions()));
TF_ASSERT_OK(original_session->Create(original_graph_def));
std::vector<Tensor> original_outputs;
TF_ASSERT_OK(original_session->Run({}, {"output"}, {}, &original_outputs));
GraphDef fused_graph_def;
TF_ASSERT_OK(FoldOldBatchNorms(original_graph_def, {{}, {"output"}},
&fused_graph_def));
std::unique_ptr<Session> fused_session(NewSession(SessionOptions()));
TF_ASSERT_OK(fused_session->Create(fused_graph_def));
std::vector<Tensor> fused_outputs;
TF_ASSERT_OK(fused_session->Run({}, {"output"}, {}, &fused_outputs));
test::ExpectTensorNear<float>(original_outputs[0], fused_outputs[0], 1e-5);
for (const NodeDef& node : fused_graph_def.node()) {
EXPECT_NE("FusedBatchNormWithBatchToSpace", node.op());
}
}
TEST_F(FoldOldBatchNormsTest, TestFoldOldBatchNorms) {
TestFoldOldBatchNorms();
}
TEST_F(FoldOldBatchNormsTest, TestFoldFusedBatchNorms) {
TestFoldFusedBatchNorms();
}
TEST_F(FoldOldBatchNormsTest, TestFoldFusedBatchNormsWithConcat) {
// Test axis is not 3, so all weights and offsets are fused to each of inputs
// of conv2d.
TestFoldFusedBatchNormsWithConcat(/*split=*/true);
// Test axis = 3, BatchNorm weights and offsets will be split before fused
// with conv2d weights.
TestFoldFusedBatchNormsWithConcat(/*split=*/false);
}
TEST_F(FoldOldBatchNormsTest, TestFoldFusedBatchNormsWithBatchToSpace) {
TestFoldFusedBatchNormsWithBatchToSpace();
}
TEST_F(FoldOldBatchNormsTest, TestFoldOldBatchNormsAfterDepthwiseConv2dNative) {
TestFoldOldBatchNormsAfterDepthwiseConv2dNative();
}
TEST_F(FoldOldBatchNormsTest,
TestFoldFusedBatchNormsAfterDepthwiseConv2dNative) {
TestFoldFusedBatchNormsAfterDepthwiseConv2dNative();
}
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/freeze_requantization_ranges.cc
+213
View File
@@ -0,0 +1,213 @@
/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/lib/strings/str_util.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
struct MinMaxRecord {
std::string name;
float min;
float max;
};
// Try to parse a log file containing loosely-structured lines, some of which
// are the min/max logs we want.
absl::Status ExtractMinMaxRecords(const std::string& log_file_name,
std::vector<MinMaxRecord>* records) {
std::string file_data;
TF_RETURN_IF_ERROR(
ReadFileToString(Env::Default(), log_file_name, &file_data));
const std::string print_suffix("__print__");
const std::string requant_prefix("__requant_min_max:");
std::vector<std::string> file_lines = str_util::Split(file_data, '\n');
for (const std::string& file_line : file_lines) {
// We expect to find a line with components separated by semicolons, so to
// start make sure that the basic structure is in place/
if (!absl::StrContains(file_line, print_suffix + ";" + requant_prefix)) {
continue;
}
std::vector<std::string> line_parts = str_util::Split(file_line, ';');
if (line_parts.size() < 2) {
continue;
}
// Now we want to figure out which components have the name and min max
// values by scanning for the prefix we expect.
bool min_max_found = false;
int min_max_index;
for (int i = 1; i < line_parts.size(); ++i) {
if (absl::StartsWith(line_parts[i], requant_prefix)) {
min_max_found = true;
min_max_index = i;
}
}
if (!min_max_found) {
continue;
}
// Finally we need to break out the values from the strings, and parse them
// into a form we can use.
std::string min_max_string = line_parts[min_max_index];
std::vector<std::string> min_max_parts =
str_util::Split(min_max_string, '[');
if ((min_max_parts.size() != 3) || (min_max_parts[0] != requant_prefix)) {
continue;
}
std::string min_string = min_max_parts[1];
std::vector<std::string> min_string_parts =
str_util::Split(min_string, ']');
if (min_string_parts.size() != 2) {
continue;
}
std::string min_number_string = min_string_parts[0];
float min;
if (!absl::SimpleAtof(min_number_string.c_str(), &min)) {
continue;
}
std::string max_string = min_max_parts[2];
std::vector<std::string> max_string_parts =
str_util::Split(max_string, ']');
if (max_string_parts.size() != 2) {
continue;
}
std::string max_number_string = max_string_parts[0];
float max;
if (!absl::SimpleAtof(max_number_string.c_str(), &max)) {
continue;
}
absl::string_view name_string = line_parts[min_max_index - 1];
if (!absl::EndsWith(name_string, print_suffix)) {
continue;
}
std::string name(
name_string.substr(0, name_string.size() - print_suffix.size()));
records->push_back({name, min, max});
}
return absl::OkStatus();
}
// Uses the observed min/max values for requantization captured in a log file to
// replace costly RequantizationRange ops with simple Consts.
absl::Status FreezeRequantizationRanges(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def) {
std::string min_max_log_file;
TF_RETURN_IF_ERROR(
context.GetOneStringParameter("min_max_log_file", "", &min_max_log_file));
if (min_max_log_file.empty()) {
return absl::InvalidArgumentError(
"You must pass a file name to min_max_log_file");
}
float min_percentile;
TF_RETURN_IF_ERROR(
context.GetOneFloatParameter("min_percentile", 5.0f, &min_percentile));
float max_percentile;
TF_RETURN_IF_ERROR(
context.GetOneFloatParameter("max_percentile", 5.0f, &max_percentile));
std::vector<MinMaxRecord> records;
TF_RETURN_IF_ERROR(ExtractMinMaxRecords(min_max_log_file, &records));
if (records.empty()) {
return absl::InvalidArgumentError(
"No min/max range logs were found in the log file");
}
std::map<std::string, const NodeDef*> node_map;
MapNamesToNodes(input_graph_def, &node_map);
bool any_missing_nodes = false;
std::map<std::string, std::vector<MinMaxRecord>> records_by_node;
for (const MinMaxRecord& record : records) {
records_by_node[record.name].push_back(record);
if (!node_map.count(record.name)) {
any_missing_nodes = true;
LOG(WARNING) << "Node from log not found in graph: " << record.name;
}
}
if (any_missing_nodes) {
return absl::InvalidArgumentError(
"Nodes were found in the log file that aren't present in the graph");
}
// Now find out the largest and smallest min/max values for the node.
std::map<std::string, std::pair<float, float>> range_for_nodes;
for (const auto& record_info : records_by_node) {
const std::string& name = record_info.first;
const std::vector<MinMaxRecord> records = record_info.second;
std::vector<float> mins;
std::vector<float> maxs;
for (const MinMaxRecord& record : records) {
mins.push_back(record.min);
maxs.push_back(record.max);
}
std::sort(mins.begin(), mins.end());
std::sort(maxs.begin(), maxs.end());
int min_index = std::round(mins.size() * (min_percentile / 100.0f));
if (min_index < 0) {
min_index = 0;
}
int max_index =
std::round(maxs.size() * (1.0f - (max_percentile / 100.0f)));
if (max_index > (maxs.size() - 1)) {
max_index = maxs.size() - 1;
}
const float min = mins[min_index];
const float max = maxs[max_index];
range_for_nodes[name] = {min, max};
}
std::map<std::string, std::string> inputs_to_rename;
GraphDef frozen_graph_def;
for (const NodeDef& node : input_graph_def.node()) {
if (range_for_nodes.count(node.name())) {
if (node.op() != "RequantizationRange") {
return absl::InvalidArgumentError(absl::StrCat(
"Node is expected to be a RequantizationRange op: ", node.name(),
", but is: ", node.op()));
}
const float min_value = range_for_nodes.at(node.name()).first;
NodeDef* min_node = frozen_graph_def.mutable_node()->Add();
min_node->set_op("Const");
min_node->set_name(node.name() + "/frozen_min");
SetNodeAttr("dtype", DT_FLOAT, min_node);
Tensor min_tensor(DT_FLOAT, {});
min_tensor.flat<float>()(0) = min_value;
SetNodeTensorAttr<float>("value", min_tensor, min_node);
inputs_to_rename[node.name() + ":0"] = min_node->name() + ":0";
const float max_value = range_for_nodes.at(node.name()).second;
NodeDef* max_node = frozen_graph_def.mutable_node()->Add();
max_node->set_op("Const");
max_node->set_name(node.name() + "/frozen_max");
SetNodeAttr("dtype", DT_FLOAT, max_node);
Tensor max_tensor(DT_FLOAT, {});
max_tensor.flat<float>()(0) = max_value;
SetNodeTensorAttr<float>("value", max_tensor, max_node);
inputs_to_rename[node.name() + ":1"] = max_node->name() + ":0";
} else {
NodeDef* new_node = frozen_graph_def.mutable_node()->Add();
*new_node = node;
}
}
return RenameNodeInputs(frozen_graph_def, inputs_to_rename,
std::unordered_set<std::string>(), output_graph_def);
}
REGISTER_GRAPH_TRANSFORM("freeze_requantization_ranges",
FreezeRequantizationRanges);
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/freeze_requantization_ranges_test.cc
+200
View File
@@ -0,0 +1,200 @@
/* Copyright 2015 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/cc/ops/const_op.h"
#include "tensorflow/cc/ops/image_ops.h"
#include "tensorflow/cc/ops/nn_ops.h"
#include "tensorflow/cc/ops/sendrecv_ops.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/lib/io/path.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
// Declare here, so we don't need a public header.
absl::Status FreezeRequantizationRanges(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def);
struct MinMaxRecord {
std::string name;
float min;
float max;
};
absl::Status ExtractMinMaxRecords(const std::string& log_file_name,
std::vector<MinMaxRecord>* records);
class FreezeRequantizationRangesTest : public ::testing::Test {
protected:
void TestFreezeRequantizationRanges() {
auto root = tensorflow::Scope::NewRootScope();
using namespace ::tensorflow::ops; // NOLINT(build/namespaces)
Tensor quantized_tensor(DT_QUINT8, TensorShape({1, 6}));
test::FillValues<quint8>(&quantized_tensor, {0, 0, 0, 0, 0, 0});
Output quantized_op = Const(root.WithOpName("quantized_op"),
Input::Initializer(quantized_tensor));
Tensor quantized_min_tensor(DT_FLOAT, TensorShape({}));
test::FillValues<float>(&quantized_min_tensor, {2.0f});
Output quantized_min_op = Const(root.WithOpName("quantized_min_op"),
Input::Initializer(quantized_min_tensor));
Tensor quantized_max_tensor(DT_FLOAT, TensorShape({}));
test::FillValues<float>(&quantized_max_tensor, {2.0f});
Output quantized_max_op = Const(root.WithOpName("quantized_max_op"),
Input::Initializer(quantized_min_tensor));
Tensor offset_tensor(DT_QUINT8, TensorShape({6}));
test::FillValues<quint8>(&offset_tensor, {1, 2, 3, 4, 5, 6});
Output offset_op =
Const(root.WithOpName("offset_op"), Input::Initializer(offset_tensor));
Tensor offset_min_tensor(DT_FLOAT, TensorShape({}));
test::FillValues<float>(&offset_min_tensor, {0.0f});
Output offset_min_op = Const(root.WithOpName("offset_min_op"),
Input::Initializer(offset_min_tensor));
Tensor offset_max_tensor(DT_FLOAT, TensorShape({}));
test::FillValues<float>(&offset_max_tensor, {255.0f});
Output offset_max_op = Const(root.WithOpName("offset_max_op"),
Input::Initializer(offset_max_tensor));
QuantizedBiasAdd quantized_bias_add_op(
root.WithOpName("bias_add_op"), quantized_op, offset_op,
quantized_min_op, quantized_max_op, offset_min_op, offset_max_op,
DT_QINT32);
RequantizationRange requantization_range_op(
root.WithOpName("requantization_range_op"),
quantized_bias_add_op.output, quantized_bias_add_op.min_out,
quantized_bias_add_op.max_out);
Requantize requantize_op(
root.WithOpName("requantize_op"), quantized_bias_add_op.output,
quantized_bias_add_op.min_out, quantized_bias_add_op.max_out,
requantization_range_op.output_min, requantization_range_op.output_max,
DT_QUINT8);
Output dequantize_op =
Dequantize(root.WithOpName("dequantize_op"), requantize_op.output,
requantize_op.output_min, requantize_op.output_max);
GraphDef graph_def;
TF_ASSERT_OK(root.ToGraphDef(&graph_def));
const std::string min_max_log_file_name =
io::JoinPath(testing::TmpDir(), "min_max_log_file.txt");
{
std::unique_ptr<WritableFile> file;
TF_ASSERT_OK(
Env::Default()->NewWritableFile(min_max_log_file_name, &file));
TF_ASSERT_OK(file->Append("Something irrelevant\n"));
TF_ASSERT_OK(
file->Append("[SomePrefix] "
";requantization_range_op__print__;__requant_min_max:"
"[-2.4313571][10.584145]\n"));
TF_ASSERT_OK(file->Append("Something else irrelevant\n"));
}
TransformFuncContext context;
context.input_names = {};
context.output_names = {"dequantize_op"};
context.params = {{"min_max_log_file", {min_max_log_file_name}}};
GraphDef frozen_graph_def;
TF_EXPECT_OK(
FreezeRequantizationRanges(graph_def, context, &frozen_graph_def));
std::map<std::string, const NodeDef*> node_map;
MapNamesToNodes(frozen_graph_def, &node_map);
EXPECT_EQ(0, node_map.count("requantization_range_op"));
EXPECT_EQ(1, node_map.count("requantize_op"));
const std::string& min_input =
NodeNameFromInput(node_map.at("requantize_op")->input(3));
ASSERT_EQ(1, node_map.count(min_input));
EXPECT_EQ("Const", node_map.at(min_input)->op());
const std::string& max_input =
NodeNameFromInput(node_map.at("requantize_op")->input(4));
ASSERT_EQ(1, node_map.count(max_input));
EXPECT_EQ("Const", node_map.at(max_input)->op());
std::unique_ptr<Session> original_session(NewSession(SessionOptions()));
TF_ASSERT_OK(original_session->Create(graph_def));
std::vector<Tensor> original_outputs;
TF_ASSERT_OK(
original_session->Run({}, {"dequantize_op"}, {}, &original_outputs));
std::unique_ptr<Session> frozen_session(NewSession(SessionOptions()));
TF_ASSERT_OK(frozen_session->Create(frozen_graph_def));
std::vector<Tensor> frozen_outputs;
TF_ASSERT_OK(
frozen_session->Run({}, {"dequantize_op"}, {}, &frozen_outputs));
ASSERT_EQ(original_outputs.size(), frozen_outputs.size());
ASSERT_EQ(1, frozen_outputs.size());
test::ExpectTensorNear<float>(original_outputs[0], frozen_outputs[0], 0.5);
}
void TestExtractMinMaxRecords() {
const std::string min_max_log_file_name =
io::JoinPath(testing::TmpDir(), "min_max_log_file2.txt");
{
std::unique_ptr<WritableFile> file;
TF_ASSERT_OK(
Env::Default()->NewWritableFile(min_max_log_file_name, &file));
TF_ASSERT_OK(file->Append("Something irrelevant\n"));
TF_ASSERT_OK(
file->Append("[SomePrefix] "
";requantization_range_op__print__;__requant_min_max:"
"[-2.4313571][10.584145]\n"));
TF_ASSERT_OK(file->Append("Something else irrelevant\n"));
TF_ASSERT_OK(file->Append(
"[SomeOtherPrefix] "
";other_requantization_range_op__print__;__requant_min_max:"
"[-1.0][2.0]\n"));
TF_ASSERT_OK(file->Append("Something else irrelevant\n"));
TF_ASSERT_OK(
file->Append("[SomePrefix] "
";requantization_range_op__print__;__requant_min_max:"
"[-1.bad][2.0]\n"));
}
std::vector<MinMaxRecord> records;
TF_ASSERT_OK(ExtractMinMaxRecords(min_max_log_file_name, &records));
ASSERT_EQ(2, records.size());
EXPECT_EQ("requantization_range_op", records[0].name);
EXPECT_NEAR(-2.4313571f, records[0].min, 1e-5f);
EXPECT_NEAR(10.584145f, records[0].max, 1e-5f);
EXPECT_EQ("other_requantization_range_op", records[1].name);
EXPECT_NEAR(-1.0f, records[1].min, 1e-5f);
EXPECT_NEAR(2.0f, records[1].max, 1e-5f);
}
};
TEST_F(FreezeRequantizationRangesTest, TestFreezeRequantizationRanges) {
TestFreezeRequantizationRanges();
}
TEST_F(FreezeRequantizationRangesTest, TestExtractMinMaxRecords) {
TestExtractMinMaxRecords();
}
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/fuse_convolutions.cc
+198
View File
@@ -0,0 +1,198 @@
/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/core/common_runtime/constant_folding.h"
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/graph/node_builder.h"
#include "tensorflow/core/graph/subgraph.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/tools/graph_transforms/fold_constants_lib.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
absl::Status FuseResizePadAndConv(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def) {
GraphDef replaced_graph_def;
TF_RETURN_IF_ERROR(ReplaceMatchingOpTypes(
input_graph_def, // clang-format off
{"Conv2D",
{
{"MirrorPad",
{
{"ResizeBilinear"},
{"*"}
}
},
{"*"}
}
}, // clang-format on
[](const NodeMatch& match, const std::set<std::string>& input_nodes,
const std::set<std::string>& output_nodes,
std::vector<NodeDef>* new_nodes) {
// Find all the nodes we expect in the subgraph.
const NodeDef& conv_node = match.node;
const NodeDef& mirror_pad_node = match.inputs[0].node;
const NodeDef& weights_node = match.inputs[1].node;
const NodeDef& resize_node = match.inputs[0].inputs[0].node;
const NodeDef& pad_dims_node = match.inputs[0].inputs[1].node;
// We'll be reusing the old weights and pad dimensions.
new_nodes->push_back(weights_node);
new_nodes->push_back(pad_dims_node);
// Set up the new fused version of the convolution op.
NodeDef fused_conv;
fused_conv.set_op("FusedResizeAndPadConv2D");
fused_conv.set_name(match.node.name());
AddNodeInput(resize_node.input(0), &fused_conv);
AddNodeInput(resize_node.input(1), &fused_conv);
AddNodeInput(mirror_pad_node.input(1), &fused_conv);
AddNodeInput(conv_node.input(1), &fused_conv);
CopyNodeAttr(resize_node, "align_corners", "resize_align_corners",
&fused_conv);
CopyNodeAttr(mirror_pad_node, "mode", "mode", &fused_conv);
CopyNodeAttr(conv_node, "T", "T", &fused_conv);
CopyNodeAttr(conv_node, "padding", "padding", &fused_conv);
CopyNodeAttr(conv_node, "strides", "strides", &fused_conv);
new_nodes->push_back(fused_conv);
return absl::OkStatus();
},
{}, &replaced_graph_def));
*output_graph_def = replaced_graph_def;
return absl::OkStatus();
}
absl::Status FuseResizeAndConv(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def) {
GraphDef replaced_graph_def;
TF_RETURN_IF_ERROR(ReplaceMatchingOpTypes(
input_graph_def, // clang-format off
{"Conv2D",
{
{"ResizeBilinear"},
{"*"}
}
}, // clang-format on
[](const NodeMatch& match, const std::set<std::string>& input_nodes,
const std::set<std::string>& output_nodes,
std::vector<NodeDef>* new_nodes) {
// Find all the nodes we expect in the subgraph.
const NodeDef& conv_node = match.node;
const NodeDef& resize_node = match.inputs[0].node;
const NodeDef& weights_node = match.inputs[1].node;
// We'll be reusing the old weights.
new_nodes->push_back(weights_node);
// Create a 'no-op' mirror padding node that has no effect.
NodeDef pad_dims_node;
pad_dims_node.set_op("Const");
pad_dims_node.set_name(conv_node.name() + "_dummy_paddings");
SetNodeAttr("dtype", DT_INT32, &pad_dims_node);
SetNodeTensorAttr<int32_t>("value", {4, 2}, {0, 0, 0, 0, 0, 0, 0, 0},
&pad_dims_node);
new_nodes->push_back(pad_dims_node);
// Set up the new fused version of the convolution op.
NodeDef fused_conv;
fused_conv.set_op("FusedResizeAndPadConv2D");
fused_conv.set_name(match.node.name());
AddNodeInput(resize_node.input(0), &fused_conv);
AddNodeInput(resize_node.input(1), &fused_conv);
AddNodeInput(pad_dims_node.name(), &fused_conv);
AddNodeInput(conv_node.input(1), &fused_conv);
CopyNodeAttr(resize_node, "align_corners", "resize_align_corners",
&fused_conv);
SetNodeAttr("mode", "REFLECT", &fused_conv);
CopyNodeAttr(conv_node, "T", "T", &fused_conv);
CopyNodeAttr(conv_node, "padding", "padding", &fused_conv);
CopyNodeAttr(conv_node, "strides", "strides", &fused_conv);
new_nodes->push_back(fused_conv);
return absl::OkStatus();
},
{}, &replaced_graph_def));
*output_graph_def = replaced_graph_def;
return absl::OkStatus();
}
absl::Status FusePadAndConv(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def) {
GraphDef replaced_graph_def;
TF_RETURN_IF_ERROR(ReplaceMatchingOpTypes(
input_graph_def, // clang-format off
{"Conv2D",
{
{"MirrorPad",
{
{"*"},
{"*"},
}
},
{"*"}
}
}, // clang-format on
[](const NodeMatch& match, const std::set<std::string>& input_nodes,
const std::set<std::string>& output_nodes,
std::vector<NodeDef>* new_nodes) {
// Find all the nodes we expect in the subgraph.
const NodeDef& conv_node = match.node;
CHECK_EQ("Conv2D", conv_node.op());
const NodeDef& mirror_pad_node = match.inputs[0].node;
CHECK_EQ("MirrorPad", mirror_pad_node.op());
const NodeDef& weights_node = match.inputs[1].node;
const NodeDef& input_node = match.inputs[0].inputs[0].node;
const NodeDef& pad_dims_node = match.inputs[0].inputs[1].node;
// We'll be reusing the old weights and pad dimensions.
new_nodes->push_back(weights_node);
new_nodes->push_back(input_node);
new_nodes->push_back(pad_dims_node);
// Set up the new fused version of the convolution op.
NodeDef fused_conv;
fused_conv.set_op("FusedPadConv2D");
fused_conv.set_name(match.node.name());
AddNodeInput(mirror_pad_node.input(0), &fused_conv);
AddNodeInput(mirror_pad_node.input(1), &fused_conv);
AddNodeInput(conv_node.input(1), &fused_conv);
CopyNodeAttr(mirror_pad_node, "mode", "mode", &fused_conv);
CopyNodeAttr(conv_node, "T", "T", &fused_conv);
CopyNodeAttr(conv_node, "padding", "padding", &fused_conv);
CopyNodeAttr(conv_node, "strides", "strides", &fused_conv);
new_nodes->push_back(fused_conv);
return absl::OkStatus();
},
{}, &replaced_graph_def));
*output_graph_def = replaced_graph_def;
return absl::OkStatus();
}
REGISTER_GRAPH_TRANSFORM("fuse_resize_pad_and_conv", FuseResizePadAndConv);
REGISTER_GRAPH_TRANSFORM("fuse_resize_and_conv", FuseResizeAndConv);
REGISTER_GRAPH_TRANSFORM("fuse_pad_and_conv", FusePadAndConv);
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/fuse_convolutions_test.cc
+212
View File
@@ -0,0 +1,212 @@
/* Copyright 2015 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/cc/ops/const_op.h"
#include "tensorflow/cc/ops/image_ops.h"
#include "tensorflow/cc/ops/nn_ops.h"
#include "tensorflow/cc/ops/sendrecv_ops.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
// Declare here, so we don't need a public header.
absl::Status FuseResizePadAndConv(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def);
absl::Status FuseResizeAndConv(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def);
absl::Status FusePadAndConv(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def);
class FuseConvolutionsTest : public ::testing::Test {
protected:
void TestFuseResizePadAndConv() {
auto root = tensorflow::Scope::NewRootScope();
using namespace ::tensorflow::ops; // NOLINT(build/namespaces)
Tensor input_data(DT_FLOAT, TensorShape({1, 2, 3, 2}));
test::FillValues<float>(
&input_data, {1.0f, 4.0f, 2.0f, 5.0f, 3.0f, 6.0f, -1.0f, -4.0f, -2.0f,
-5.0f, -3.0f, -6.0f});
Output input_op =
Const(root.WithOpName("input_op"), Input::Initializer(input_data));
Output resize_op = ResizeBilinear(root.WithOpName("resize_op"), input_op,
Const(root.WithOpName("size"), {12, 4}),
ResizeBilinear::AlignCorners(false));
Tensor pad_dims_data(DT_INT32, TensorShape({4, 2}));
test::FillValues<int32_t>(&pad_dims_data, {0, 0, 1, 1, 2, 2, 0, 0});
Output pad_dims_op = Const(root.WithOpName("pad_dims_op"),
Input::Initializer(pad_dims_data));
Output pad_op =
MirrorPad(root.WithOpName("pad_op"), resize_op, pad_dims_op, "REFLECT");
Tensor weights_data(DT_FLOAT, TensorShape({1, 2, 2, 2}));
test::FillValues<float>(&weights_data,
{1.0f, 2.0f, 3.0f, 4.0f, 0.1f, 0.2f, 0.3f, 0.4f});
Output weights_op =
Const(root.WithOpName("weights_op"), Input::Initializer(weights_data));
Output conv_op = Conv2D(root.WithOpName("output"), pad_op, weights_op,
{1, 1, 1, 1}, "VALID");
GraphDef original_graph_def;
TF_ASSERT_OK(root.ToGraphDef(&original_graph_def));
std::unique_ptr<Session> original_session(NewSession(SessionOptions()));
TF_ASSERT_OK(original_session->Create(original_graph_def));
std::vector<Tensor> original_outputs;
TF_ASSERT_OK(original_session->Run({}, {"output"}, {}, &original_outputs));
GraphDef fused_graph_def;
TF_ASSERT_OK(FuseResizePadAndConv(original_graph_def, {{}, {"output"}},
&fused_graph_def));
std::unique_ptr<Session> fused_session(NewSession(SessionOptions()));
TF_ASSERT_OK(fused_session->Create(fused_graph_def));
std::vector<Tensor> fused_outputs;
TF_ASSERT_OK(fused_session->Run({}, {"output"}, {}, &fused_outputs));
test::ExpectTensorNear<float>(original_outputs[0], fused_outputs[0], 1e-5);
for (const NodeDef& node : fused_graph_def.node()) {
EXPECT_NE("Conv2D", node.op());
EXPECT_NE("MirrorPad", node.op());
EXPECT_NE("ResizeBilinear", node.op());
}
}
void TestFuseResizeAndConv() {
auto root = tensorflow::Scope::NewRootScope();
using namespace ::tensorflow::ops; // NOLINT(build/namespaces)
Tensor input_data(DT_FLOAT, TensorShape({1, 2, 3, 2}));
test::FillValues<float>(
&input_data, {1.0f, 4.0f, 2.0f, 5.0f, 3.0f, 6.0f, -1.0f, -4.0f, -2.0f,
-5.0f, -3.0f, -6.0f});
Output input_op =
Const(root.WithOpName("input_op"), Input::Initializer(input_data));
Output resize_op = ResizeBilinear(root.WithOpName("resize_op"), input_op,
Const(root.WithOpName("size"), {12, 4}),
ResizeBilinear::AlignCorners(false));
Tensor weights_data(DT_FLOAT, TensorShape({1, 2, 2, 2}));
test::FillValues<float>(&weights_data,
{1.0f, 2.0f, 3.0f, 4.0f, 0.1f, 0.2f, 0.3f, 0.4f});
Output weights_op =
Const(root.WithOpName("weights_op"), Input::Initializer(weights_data));
Output conv_op = Conv2D(root.WithOpName("output"), resize_op, weights_op,
{1, 1, 1, 1}, "VALID");
GraphDef original_graph_def;
TF_ASSERT_OK(root.ToGraphDef(&original_graph_def));
std::unique_ptr<Session> original_session(NewSession(SessionOptions()));
TF_ASSERT_OK(original_session->Create(original_graph_def));
std::vector<Tensor> original_outputs;
TF_ASSERT_OK(original_session->Run({}, {"output"}, {}, &original_outputs));
GraphDef fused_graph_def;
TF_ASSERT_OK(FuseResizeAndConv(original_graph_def, {{}, {"output"}},
&fused_graph_def));
std::unique_ptr<Session> fused_session(NewSession(SessionOptions()));
TF_ASSERT_OK(fused_session->Create(fused_graph_def));
std::vector<Tensor> fused_outputs;
TF_ASSERT_OK(fused_session->Run({}, {"output"}, {}, &fused_outputs));
test::ExpectTensorNear<float>(original_outputs[0], fused_outputs[0], 1e-5);
for (const NodeDef& node : fused_graph_def.node()) {
EXPECT_NE("Conv2D", node.op());
EXPECT_NE("ResizeBilinear", node.op());
}
}
void TestFusePadAndConv() {
auto root = tensorflow::Scope::NewRootScope();
using namespace ::tensorflow::ops; // NOLINT(build/namespaces)
Tensor input_data(DT_FLOAT, TensorShape({1, 2, 3, 2}));
test::FillValues<float>(
&input_data, {1.0f, 4.0f, 2.0f, 5.0f, 3.0f, 6.0f, -1.0f, -4.0f, -2.0f,
-5.0f, -3.0f, -6.0f});
Output input_op =
Const(root.WithOpName("input_op"), Input::Initializer(input_data));
Tensor pad_dims_data(DT_INT32, TensorShape({4, 2}));
test::FillValues<int32_t>(&pad_dims_data, {0, 0, 1, 1, 2, 2, 0, 0});
Output pad_dims_op = Const(root.WithOpName("pad_dims_op"),
Input::Initializer(pad_dims_data));
Output pad_op =
MirrorPad(root.WithOpName("pad_op"), input_op, pad_dims_op, "REFLECT");
Tensor weights_data(DT_FLOAT, TensorShape({1, 2, 2, 2}));
test::FillValues<float>(&weights_data,
{1.0f, 2.0f, 3.0f, 4.0f, 0.1f, 0.2f, 0.3f, 0.4f});
Output weights_op =
Const(root.WithOpName("weights_op"), Input::Initializer(weights_data));
Output conv_op = Conv2D(root.WithOpName("output"), pad_op, weights_op,
{1, 1, 1, 1}, "VALID");
GraphDef original_graph_def;
TF_ASSERT_OK(root.ToGraphDef(&original_graph_def));
std::unique_ptr<Session> original_session(NewSession(SessionOptions()));
TF_ASSERT_OK(original_session->Create(original_graph_def));
std::vector<Tensor> original_outputs;
TF_ASSERT_OK(original_session->Run({}, {"output"}, {}, &original_outputs));
GraphDef fused_graph_def;
TF_ASSERT_OK(
FusePadAndConv(original_graph_def, {{}, {"output"}}, &fused_graph_def));
std::unique_ptr<Session> fused_session(NewSession(SessionOptions()));
TF_ASSERT_OK(fused_session->Create(fused_graph_def));
std::vector<Tensor> fused_outputs;
TF_ASSERT_OK(fused_session->Run({}, {"output"}, {}, &fused_outputs));
test::ExpectTensorNear<float>(original_outputs[0], fused_outputs[0], 1e-5);
for (const NodeDef& node : fused_graph_def.node()) {
EXPECT_NE("Conv2D", node.op());
EXPECT_NE("MirrorPad", node.op());
}
}
};
TEST_F(FuseConvolutionsTest, TestFuseResizePadAndConv) {
TestFuseResizePadAndConv();
}
TEST_F(FuseConvolutionsTest, TestFuseResizeAndConv) { TestFuseResizeAndConv(); }
TEST_F(FuseConvolutionsTest, TestFusePadAndConv) { TestFusePadAndConv(); }
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/inline_partitionedcall.cc
+151
View File
@@ -0,0 +1,151 @@
/* 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 <string>
#include "absl/algorithm/container.h"
#include "absl/container/flat_hash_map.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/op_def.pb.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
constexpr char kPartitionedCallOpName[] = "PartitionedCall";
constexpr char kFunctionAttrName[] = "f";
namespace {
absl::optional<FunctionDef> GetFunctionByNameFromLibrary(
const GraphDef& graph, absl::string_view function_name) {
for (const auto& fct : graph.library().function()) {
if (fct.signature().name() == function_name) {
return fct;
}
}
return {};
}
std::string NormalizeNodeDefInput(const std::string& input_name) {
std::vector<std::string> name_parts =
absl::StrSplit(input_name, absl::ByChar(':'));
if (name_parts.size() > 2) {
return absl::StrCat(name_parts[0], ":", name_parts.back());
}
return input_name;
}
} // namespace
absl::Status InlinePartitionedCall(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def) {
output_graph_def->Clear();
absl::flat_hash_map<std::string, std::string> remap_input;
for (const NodeDef& node : input_graph_def.node()) {
if (node.op() == kPartitionedCallOpName) {
if (node.attr().count(kFunctionAttrName) == 0) {
return absl::Status(
absl::StatusCode::kNotFound,
"Node " + node.name() + " has no attribute: " + kFunctionAttrName);
}
if (!node.attr().at(kFunctionAttrName).has_func()) {
return absl::Status(absl::StatusCode::kNotFound,
"Cannot figure out function name");
}
const std::string function_name =
node.attr().at(kFunctionAttrName).func().name();
absl::optional<FunctionDef> function =
GetFunctionByNameFromLibrary(input_graph_def, function_name);
if (!function.has_value()) {
return absl::Status(absl::StatusCode::kNotFound,
"function " + function_name + " Not found");
}
const std::string prefix = node.name();
const int kOutputArgumentCount =
function->signature().output_arg().size();
for (int k = 0; k < kOutputArgumentCount; ++k) {
const std::string function_arg_output_name =
function->ret().at(function->signature().output_arg()[k].name());
remap_input.insert_or_assign(
CanonicalInputName(absl::StrCat(node.name(), ":", k)),
absl::StrCat(prefix, "/",
NormalizeNodeDefInput(function_arg_output_name)));
}
const int kInputArgumentCount = function->signature().input_arg().size();
if (node.input().size() != kInputArgumentCount) {
return absl::Status(absl::StatusCode::kInvalidArgument,
"Called function " + function_name +
" has invalid input signature.");
}
absl::flat_hash_map<std::string, std::string> input_argument_map;
for (int k = 0; k < kInputArgumentCount; ++k) {
const std::string canonical_name =
CanonicalInputName(function->signature().input_arg()[k].name());
input_argument_map.insert_or_assign(canonical_name, node.input()[k]);
}
for (const NodeDef& function_node : function->node_def()) {
NodeDef* new_node = output_graph_def->mutable_node()->Add();
*new_node = function_node;
new_node->set_name(absl::StrCat(prefix, "/", function_node.name()));
absl::c_transform(
*new_node->mutable_input(), new_node->mutable_input()->begin(),
[prefix, input_argument_map](const std::string& input_name) {
const std::string canonical_input_name =
CanonicalInputName(input_name);
if (input_argument_map.find(canonical_input_name) !=
input_argument_map.end()) {
return input_argument_map.at(canonical_input_name);
}
return absl::StrCat(prefix, "/",
NormalizeNodeDefInput(input_name));
});
}
} else {
NodeDef* new_node = output_graph_def->mutable_node()->Add();
*new_node = node;
}
}
// Remap PartitionCall outputs to correct nodes.
for (NodeDef& node : *output_graph_def->mutable_node()) {
absl::c_transform(
*node.mutable_input(), node.mutable_input()->begin(),
[remap_input](const std::string& input_name) {
const std::string canonical_input_name =
CanonicalInputName(input_name);
if (remap_input.find(canonical_input_name) != remap_input.end()) {
return remap_input.at(canonical_input_name);
}
return input_name;
});
}
return absl::OkStatus();
}
REGISTER_GRAPH_TRANSFORM("inline_partitionedcall", InlinePartitionedCall);
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/inline_partitionedcall_test.cc
+136
View File
@@ -0,0 +1,136 @@
/* 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 <string>
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/protobuf.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
constexpr char kGraphDefWithPartitionedCall[] =
"node {\n"
" name: \"y\"\n"
" op: \"Placeholder\"\n"
"}\n"
"node {\n"
" name: \"sub/y\"\n"
" op: \"Const\"\n"
"}\n"
"node {\n"
" name: \"PartitionedCall\"\n"
" op: \"PartitionedCall\"\n"
" input: \"y\"\n"
" input: \"sub/y\"\n"
" attr {\n"
" key: \"f\"\n"
" value {\n"
" func {\n"
" name: \"__inference_simple_add_14\"\n"
" }\n"
" }\n"
" }\n"
"}\n"
"node {\n"
" name: \"add/y\"\n"
" op: \"Const\"\n"
"}\n"
"node {\n"
" name: \"add\"\n"
" op: \"AddV2\"\n"
" input: \"PartitionedCall\"\n"
" input: \"add/y\"\n"
"}\n"
"node {\n"
" name: \"Identity\"\n"
" op: \"Identity\"\n"
" input: \"add\"\n"
"}\n"
"library {\n"
" function {\n"
" signature {\n"
" name: \"__inference_simple_add_14\"\n"
" input_arg {\n"
" name: \"x\"\n"
" type: DT_FLOAT\n"
" }\n"
" input_arg {\n"
" name: \"y\"\n"
" type: DT_FLOAT\n"
" }\n"
" output_arg {\n"
" name: \"identity\"\n"
" type: DT_FLOAT\n"
" }\n"
" }\n"
" node_def {\n"
" name: \"mul\"\n"
" op: \"Mul\"\n"
" input: \"x\"\n"
" input: \"y\"\n"
" }\n"
" node_def {\n"
" name: \"add/y\"\n"
" op: \"Const\"\n"
" }\n"
" node_def {\n"
" name: \"add\"\n"
" op: \"AddV2\"\n"
" input: \"mul:z:0\"\n"
" input: \"add/y:output:0\"\n"
" }\n"
" node_def {\n"
" name: \"Identity\"\n"
" op: \"Identity\"\n"
" input: \"add:z:0\"\n"
" }\n"
" ret {\n"
" key: \"identity\"\n"
" value: \"Identity:output:0\"\n"
" }\n"
" }\n"
"}\n";
// Declare here, so we don't need a public header.
absl::Status InlinePartitionedCall(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def);
TEST(InlinePartitionedCallTest, Inlining) {
GraphDef in_graph;
EXPECT_TRUE(::tensorflow::protobuf::TextFormat::ParseFromString(
kGraphDefWithPartitionedCall, &in_graph));
GraphDef result;
TransformFuncContext context;
context.input_names = {"y"};
context.output_names = {"Identity"};
TF_ASSERT_OK(InlinePartitionedCall(in_graph, context, &result));
EXPECT_TRUE(std::none_of(
result.node().cbegin(), result.node().cend(),
[](const NodeDef& node) { return node.op() == "PartitionedCall"; }));
EXPECT_EQ(9, result.node().size());
TF_EXPECT_OK(IsGraphValid(result));
}
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/insert_logging.cc
+151
View File
@@ -0,0 +1,151 @@
/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/core/common_runtime/constant_folding.h"
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/graph/node_builder.h"
#include "tensorflow/core/graph/subgraph.h"
#include "tensorflow/core/lib/strings/str_util.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/tools/graph_transforms/fold_constants_lib.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
// Clears the device field of all ops in the graph.
absl::Status InsertLogging(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def) {
std::unordered_set<std::string> ops;
bool has_ops;
if (context.params.count("op")) {
has_ops = true;
for (const std::string& op : context.params.at("op")) {
ops.insert(op);
}
} else {
has_ops = false;
}
std::unordered_set<std::string> prefixes;
bool has_prefixes;
if (context.params.count("prefix")) {
has_prefixes = true;
for (const std::string& prefix : context.params.at("prefix")) {
prefixes.insert(prefix);
}
} else {
has_prefixes = false;
}
std::string message;
TF_RETURN_IF_ERROR(context.GetOneStringParameter("message", "", &message));
bool show_name;
TF_RETURN_IF_ERROR(
context.GetOneBoolParameter("show_name", false, &show_name));
bool show_op;
TF_RETURN_IF_ERROR(context.GetOneBoolParameter("show_op", false, &show_op));
int32_t first_n;
TF_RETURN_IF_ERROR(context.GetOneInt32Parameter("first_n", -1, &first_n));
int32_t summarize;
TF_RETURN_IF_ERROR(
context.GetOneInt32Parameter("summarize", 1024, &summarize));
std::unordered_map<std::string, std::set<int>> node_outputs;
for (const NodeDef& node : input_graph_def.node()) {
for (const std::string& input : node.input()) {
const std::string canonical_input = CanonicalInputName(input);
std::string prefix;
std::string name;
std::string suffix;
NodeNamePartsFromInput(canonical_input, &prefix, &name, &suffix);
const std::string output_index_string =
suffix.substr(1, suffix.size() - 1);
int32_t output_index;
if (!absl::SimpleAtoi(output_index_string, &output_index)) {
return absl::InvalidArgumentError(
absl::StrCat("Couldn't understand output number in ", input));
}
node_outputs[name].insert(output_index);
}
}
std::map<std::string, std::string> inputs_to_rename;
std::unordered_set<std::string> ignore_when_renaming;
GraphDef logged_graph_def;
for (const NodeDef& node : input_graph_def.node()) {
NodeDef* new_node = logged_graph_def.mutable_node()->Add();
*new_node = node;
if (node_outputs[node.name()].empty()) {
// There were no outputs found to this node, so skip it.
continue;
}
const bool op_matches = (ops.count(node.op()) > 0);
bool prefix_matches = false;
for (const std::string& prefix : prefixes) {
if (absl::StartsWith(node.name(), prefix)) {
prefix_matches = true;
}
}
// If we're not looking for ops, or we found the right op, and if we're not
// looking for prefixes or we found the right prefix, then add logging here.
if ((!has_ops || op_matches) && (!has_prefixes || prefix_matches)) {
const std::string name_suffix = "__print__";
DataTypeVector input_types;
DataTypeVector output_types;
TF_RETURN_IF_ERROR(GetInOutTypes(node, &input_types, &output_types));
NodeDef* print_node = logged_graph_def.mutable_node()->Add();
print_node->set_op("Print");
print_node->set_name(absl::StrCat(node.name(), name_suffix));
std::string node_message;
if (show_op) {
node_message += ";" + node.op() + ";";
}
if (show_name) {
node_message += ";" + print_node->name() + ";";
}
node_message += message;
SetNodeAttr("message", node_message, print_node);
SetNodeAttr("first_n", first_n, print_node);
SetNodeAttr("summarize", summarize, print_node);
print_node->add_input(node.name() + ":0");
SetNodeAttr("T", output_types[0], print_node);
for (int output_index : node_outputs[node.name()]) {
print_node->add_input(absl::StrCat(node.name(), ":", output_index));
}
SetNodeAttr("U", output_types, print_node);
ignore_when_renaming.insert(print_node->name());
// Rewrite the graph so all references to the first input of the original
// op now pull from the print op instead, so it's executed.
inputs_to_rename[node.name() + ":0"] =
absl::StrCat(node.name(), name_suffix, ":0");
}
}
output_graph_def->Clear();
return RenameNodeInputs(logged_graph_def, inputs_to_rename,
ignore_when_renaming, output_graph_def);
}
REGISTER_GRAPH_TRANSFORM("insert_logging", InsertLogging);
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/insert_logging_test.cc
+203
View File
@@ -0,0 +1,203 @@
/* Copyright 2015 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/cc/ops/const_op.h"
#include "tensorflow/cc/ops/image_ops.h"
#include "tensorflow/cc/ops/nn_ops.h"
#include "tensorflow/cc/ops/sendrecv_ops.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
// Declare here, so we don't need a public header.
absl::Status InsertLogging(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def);
class InsertLoggingTest : public ::testing::Test {
protected:
void CheckGraphCanRun(const GraphDef& graph_def,
const std::vector<std::string>& output_names) {
std::unique_ptr<Session> session(NewSession(SessionOptions()));
TF_ASSERT_OK(session->Create(graph_def));
std::vector<Tensor> outputs;
TF_ASSERT_OK(session->Run({}, output_names, {}, &outputs));
}
void TestInsertLogging() {
auto root = tensorflow::Scope::NewRootScope();
using namespace ::tensorflow::ops; // NOLINT(build/namespaces)
Tensor const_tensor(DT_FLOAT, TensorShape({10}));
test::FillIota<float>(&const_tensor, 1.0f);
Output const_node1 =
Const(root.WithOpName("const_node1"), Input::Initializer(const_tensor));
Output const_node2 =
Const(root.WithOpName("const_node2"), Input::Initializer(const_tensor));
Output const_node3 =
Const(root.WithOpName("const_node3"), Input::Initializer(const_tensor));
Output add_node2 =
Add(root.WithOpName("add_node2"), const_node1, const_node2);
Output add_node3 =
Add(root.WithOpName("add_node3"), const_node1, const_node3);
Output mul_node1 = Mul(root.WithOpName("mul_node1"), add_node2, add_node3);
Output add_node4 =
Add(root.WithOpName("add_node4"), mul_node1, const_node3);
GraphDef graph_def;
TF_ASSERT_OK(root.ToGraphDef(&graph_def));
CheckGraphCanRun(graph_def, {"add_node4"});
GraphDef result;
TransformFuncContext context;
context.input_names = {};
context.output_names = {"add_node4"};
TF_ASSERT_OK(InsertLogging(graph_def, context, &result));
CheckGraphCanRun(result, {"add_node4"});
std::unordered_set<std::string> print_inputs;
for (const NodeDef& node : result.node()) {
if (node.op() == "Print") {
print_inputs.insert(node.input(0));
}
}
EXPECT_EQ(6, print_inputs.size());
EXPECT_EQ(1, print_inputs.count("mul_node1:0"));
EXPECT_EQ(1, print_inputs.count("add_node2:0"));
EXPECT_EQ(1, print_inputs.count("add_node3:0"));
EXPECT_EQ(0, print_inputs.count("add_node4:0"));
EXPECT_EQ(1, print_inputs.count("const_node1:0"));
EXPECT_EQ(1, print_inputs.count("const_node2:0"));
EXPECT_EQ(1, print_inputs.count("const_node3:0"));
}
void TestInsertLoggingByOpType() {
auto root = tensorflow::Scope::NewRootScope();
using namespace ::tensorflow::ops; // NOLINT(build/namespaces)
Tensor const_tensor(DT_FLOAT, TensorShape({10}));
test::FillIota<float>(&const_tensor, 1.0f);
Output const_node1 =
Const(root.WithOpName("const_node1"), Input::Initializer(const_tensor));
Output const_node2 =
Const(root.WithOpName("const_node2"), Input::Initializer(const_tensor));
Output const_node3 =
Const(root.WithOpName("const_node3"), Input::Initializer(const_tensor));
Output add_node2 =
Add(root.WithOpName("add_node2"), const_node1, const_node2);
Output add_node3 =
Add(root.WithOpName("add_node3"), const_node1, const_node3);
Output mul_node1 = Mul(root.WithOpName("mul_node1"), add_node2, add_node3);
Output add_node4 =
Add(root.WithOpName("add_node4"), mul_node1, const_node3);
GraphDef graph_def;
TF_ASSERT_OK(root.ToGraphDef(&graph_def));
CheckGraphCanRun(graph_def, {"add_node4"});
GraphDef result;
TransformFuncContext context;
context.input_names = {};
context.output_names = {"add_node4"};
context.params.insert(std::pair<std::string, std::vector<std::string>>(
{"op", {"Mul", "Add"}}));
TF_ASSERT_OK(InsertLogging(graph_def, context, &result));
CheckGraphCanRun(result, {"add_node4"});
std::unordered_set<std::string> print_inputs;
for (const NodeDef& node : result.node()) {
if (node.op() == "Print") {
print_inputs.insert(node.input(0));
}
}
EXPECT_EQ(3, print_inputs.size());
EXPECT_EQ(1, print_inputs.count("mul_node1:0"));
EXPECT_EQ(1, print_inputs.count("add_node2:0"));
EXPECT_EQ(1, print_inputs.count("add_node3:0"));
EXPECT_EQ(0, print_inputs.count("add_node4:0"));
EXPECT_EQ(0, print_inputs.count("const_node1:0"));
EXPECT_EQ(0, print_inputs.count("const_node2:0"));
EXPECT_EQ(0, print_inputs.count("const_node3:0"));
}
void TestInsertLoggingByPrefix() {
auto root = tensorflow::Scope::NewRootScope();
using namespace ::tensorflow::ops; // NOLINT(build/namespaces)
Tensor const_tensor(DT_FLOAT, TensorShape({10}));
test::FillIota<float>(&const_tensor, 1.0f);
Output const_node1 =
Const(root.WithOpName("const_node1"), Input::Initializer(const_tensor));
Output const_node2 =
Const(root.WithOpName("const_node2"), Input::Initializer(const_tensor));
Output const_node3 =
Const(root.WithOpName("const_node3"), Input::Initializer(const_tensor));
Output add_node2 =
Add(root.WithOpName("add_node2"), const_node1, const_node2);
Output add_node3 =
Add(root.WithOpName("add_node3"), const_node1, const_node3);
Output mul_node1 = Mul(root.WithOpName("mul_node1"), add_node2, add_node3);
Output add_node4 =
Add(root.WithOpName("add_node4"), mul_node1, const_node3);
GraphDef graph_def;
TF_ASSERT_OK(root.ToGraphDef(&graph_def));
CheckGraphCanRun(graph_def, {"add_node4"});
GraphDef result;
TransformFuncContext context;
context.input_names = {};
context.output_names = {"add_node4"};
context.params.insert(std::pair<std::string, std::vector<std::string>>(
{"prefix", {"add_node"}}));
TF_ASSERT_OK(InsertLogging(graph_def, context, &result));
CheckGraphCanRun(result, {"add_node4"});
std::unordered_set<std::string> print_inputs;
for (const NodeDef& node : result.node()) {
if (node.op() == "Print") {
print_inputs.insert(node.input(0));
}
}
EXPECT_EQ(2, print_inputs.size());
EXPECT_EQ(0, print_inputs.count("mul_node1:0"));
EXPECT_EQ(1, print_inputs.count("add_node2:0"));
EXPECT_EQ(1, print_inputs.count("add_node3:0"));
EXPECT_EQ(0, print_inputs.count("add_node4:0"));
EXPECT_EQ(0, print_inputs.count("const_node1:0"));
EXPECT_EQ(0, print_inputs.count("const_node2:0"));
EXPECT_EQ(0, print_inputs.count("const_node3:0"));
}
};
TEST_F(InsertLoggingTest, TestInsertLogging) { TestInsertLogging(); }
TEST_F(InsertLoggingTest, TestInsertLoggingByOpType) {
TestInsertLoggingByOpType();
}
TEST_F(InsertLoggingTest, TestInsertLoggingByPrefix) {
TestInsertLoggingByPrefix();
}
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/obfuscate_names.cc
+98
View File
@@ -0,0 +1,98 @@
/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/core/common_runtime/constant_folding.h"
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/graph/node_builder.h"
#include "tensorflow/core/graph/subgraph.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/tools/graph_transforms/fold_constants_lib.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
// Renames all nodes not uses as graph inputs or outputs to short numerical
// forms.
absl::Status ObfuscateNames(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def) {
std::unordered_set<std::string> required_nodes;
for (const std::string& input : context.input_names) {
required_nodes.insert(input);
}
for (const std::string& output : context.output_names) {
required_nodes.insert(output);
}
const std::string valid_chars =
"0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ";
const int64_t chars_size = valid_chars.size();
std::map<std::string, std::string> new_names;
int64_t name_index = 0;
for (const NodeDef& input_node : input_graph_def.node()) {
const std::string& old_name = input_node.name();
std::string new_name;
if (required_nodes.count(old_name)) {
new_name = old_name;
} else {
do {
int64_t remaining = name_index;
new_name = "";
while (true) {
const int64_t remainder = (remaining % chars_size);
const char current_char = valid_chars[remainder];
new_name = current_char + new_name;
remaining /= chars_size;
if (remaining <= 0) {
break;
}
}
++name_index;
} while (required_nodes.count(new_name));
}
new_names[old_name] = new_name;
}
output_graph_def->Clear();
for (const NodeDef& input_node : input_graph_def.node()) {
NodeDef* node = output_graph_def->mutable_node()->Add();
*node = input_node;
const std::string& old_name = input_node.name();
node->set_name(new_names[old_name]);
node->mutable_input()->Clear();
for (const std::string& input_name : input_node.input()) {
std::string prefix;
std::string input_node_name;
std::string suffix;
NodeNamePartsFromInput(input_name, &prefix, &input_node_name, &suffix);
if (new_names.count(input_node_name) == 0) {
return absl::InvalidArgumentError(absl::StrCat(
"No node named ", input_node_name, " for input to ", old_name));
}
std::string new_input_name = prefix + new_names[input_node_name] + suffix;
*(node->mutable_input()->Add()) = new_input_name;
}
}
return absl::OkStatus();
}
REGISTER_GRAPH_TRANSFORM("obfuscate_names", ObfuscateNames);
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/obfuscate_names_test.cc
+142
View File
@@ -0,0 +1,142 @@
/* Copyright 2015 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/cc/ops/const_op.h"
#include "tensorflow/cc/ops/image_ops.h"
#include "tensorflow/cc/ops/nn_ops.h"
#include "tensorflow/cc/ops/sendrecv_ops.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
// Declare here, so we don't need a public header.
absl::Status ObfuscateNames(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def);
class ObfuscateNamesTest : public ::testing::Test {
protected:
void TestSimpleTree() {
GraphDef graph_def;
NodeDef* add_node1 = graph_def.add_node();
add_node1->set_name("add_node1");
add_node1->set_op("Add");
add_node1->add_input("add_node2");
add_node1->add_input("add_node3");
NodeDef* add_node2 = graph_def.add_node();
add_node2->set_name("add_node2");
add_node2->set_op("Add");
add_node2->add_input("const_node1");
add_node2->add_input("const_node2");
NodeDef* add_node3 = graph_def.add_node();
add_node3->set_name("add_node3");
add_node3->set_op("Add");
add_node3->add_input("const_node3");
add_node3->add_input("const_node4");
NodeDef* const_node1 = graph_def.add_node();
const_node1->set_name("const_node1");
const_node1->set_op("Const");
NodeDef* const_node2 = graph_def.add_node();
const_node2->set_name("const_node2");
const_node2->set_op("Const");
NodeDef* const_node3 = graph_def.add_node();
const_node3->set_name("const_node3");
const_node3->set_op("Const");
NodeDef* const_node4 = graph_def.add_node();
const_node4->set_name("const_node4");
const_node4->set_op("Const");
GraphDef result;
TF_ASSERT_OK(
ObfuscateNames(graph_def, {{"const_node1"}, {"add_node1"}}, &result));
std::map<std::string, const NodeDef*> node_lookup;
MapNamesToNodes(result, &node_lookup);
EXPECT_EQ(1, node_lookup.count("add_node1"));
EXPECT_EQ(0, node_lookup.count("add_node2"));
EXPECT_EQ(0, node_lookup.count("add_node3"));
EXPECT_EQ(1, node_lookup.count("const_node1"));
EXPECT_EQ(0, node_lookup.count("const_node2"));
EXPECT_EQ(0, node_lookup.count("const_node3"));
EXPECT_EQ(0, node_lookup.count("const_node4"));
}
void TestManyNodes() {
GraphDef graph_def;
for (int i = 0; i < 1000; ++i) {
NodeDef* const_node = graph_def.add_node();
const_node->set_name(absl::StrCat("const_node", i));
const_node->set_op("Const");
}
GraphDef result;
TF_ASSERT_OK(ObfuscateNames(graph_def, {{"const_node0"}, {"const_node999"}},
&result));
std::map<std::string, const NodeDef*> node_lookup;
MapNamesToNodes(result, &node_lookup);
EXPECT_EQ(1, node_lookup.count("const_node0"));
EXPECT_EQ(0, node_lookup.count("const_node500"));
EXPECT_EQ(1, node_lookup.count("const_node999"));
}
void TestNameClashes() {
GraphDef graph_def;
for (int i = 0; i < 1000; ++i) {
NodeDef* const_node = graph_def.add_node();
const_node->set_name(absl::StrCat("1", i));
const_node->set_op("Const");
}
GraphDef result;
TF_ASSERT_OK(ObfuscateNames(graph_def, {{"10"}, {"19"}}, &result));
std::map<std::string, const NodeDef*> node_lookup;
MapNamesToNodes(result, &node_lookup);
EXPECT_EQ(1, node_lookup.count("10"));
EXPECT_EQ(1, node_lookup.count("19"));
std::unordered_set<std::string> names;
for (const NodeDef& node : result.node()) {
EXPECT_EQ(0, names.count(node.name()))
<< "Found multiple nodes with name '" << node.name() << "'";
names.insert(node.name());
}
}
};
TEST_F(ObfuscateNamesTest, TestSimpleTree) { TestSimpleTree(); }
TEST_F(ObfuscateNamesTest, TestManyNodes) { TestManyNodes(); }
TEST_F(ObfuscateNamesTest, TestNameClashes) { TestNameClashes(); }
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/python/transform_graph_test.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.
# ==============================================================================
"""Tests for StatSummarizer Python wrapper."""
from tensorflow.core.framework import attr_value_pb2
from tensorflow.core.framework import graph_pb2
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import tensor_util
from tensorflow.python.platform import test
from tensorflow.tools import graph_transforms
class TransformGraphTest(test.TestCase):
# This test constructs a graph with a relu op that's not used by the normal
# inference path, and then tests that the strip_unused transform removes it as
# expected.
def testTransformGraph(self):
input_graph_def = graph_pb2.GraphDef()
const_op1 = input_graph_def.node.add()
const_op1.op = "Const"
const_op1.name = "const_op1"
const_op1.attr["dtype"].CopyFrom(attr_value_pb2.AttrValue(
type=dtypes.float32.as_datatype_enum))
const_op1.attr["value"].CopyFrom(
attr_value_pb2.AttrValue(tensor=tensor_util.make_tensor_proto(
[1, 2], dtypes.float32, [1, 2])))
const_op2 = input_graph_def.node.add()
const_op2.op = "Const"
const_op2.name = "const_op2"
const_op2.attr["dtype"].CopyFrom(attr_value_pb2.AttrValue(
type=dtypes.float32.as_datatype_enum))
const_op2.attr["value"].CopyFrom(
attr_value_pb2.AttrValue(tensor=tensor_util.make_tensor_proto(
[3, 4], dtypes.float32, [1, 2])))
# Create an add that has two constants as inputs.
add_op = input_graph_def.node.add()
add_op.op = "Add"
add_op.attr["T"].CopyFrom(attr_value_pb2.AttrValue(
type=dtypes.float32.as_datatype_enum))
add_op.name = "add_op"
add_op.input.extend(["const_op1", "const_op2"])
# Create a relu that reads from the add.
relu_op = input_graph_def.node.add()
relu_op.op = "Relu"
relu_op.attr["T"].CopyFrom(attr_value_pb2.AttrValue(
type=dtypes.float32.as_datatype_enum))
relu_op.name = "relu_op"
relu_op.input.extend(["add_op"])
# We're specifying that add_op is the final output, and so the relu isn't
# needed.
input_names = []
output_names = ["add_op"]
transforms = ["strip_unused_nodes"]
transformed_graph_def = graph_transforms.TransformGraph(
input_graph_def, input_names, output_names, transforms)
# We expect that the relu is no longer present after running the transform.
for node in transformed_graph_def.node:
self.assertNotEqual("Relu", node.op)
if __name__ == "__main__":
test.main()
tensorflow/tools/graph_transforms/quantize_nodes.cc
+963
View File
@@ -0,0 +1,963 @@
/* Copyright 2016 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.
==============================================================================*/
#define EIGEN_USE_THREADS
#include "tensorflow/core/common_runtime/constant_folding.h"
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/common_runtime/threadpool_device.h"
#include "tensorflow/core/graph/node_builder.h"
#include "tensorflow/core/graph/subgraph.h"
#include "tensorflow/core/kernels/quantization_utils.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
// Holds the information we need to translate from a float version of this op
// into the quantized equivalent.
struct QuantizedOpInfo {
// The name of the float op.
std::string float_name;
// Which attributes to copy directly over.
std::vector<std::string> attrs_to_copy;
// Extra data type attributes we need to set.
std::vector<std::pair<std::string, DataType>> dtypes_to_set;
// What depth of inputs the op can read in.
DataType input_bit_depth;
// The depth of the op's quantized outputs.
DataType output_bit_depth;
// Which inputs (e.g. shapes) aren't involved in the quantization process.
std::set<int32_t> unquantized_inputs;
// How the outputs are arranged, either
// [input0, input1, min0, max0, min1, max1] for contiguous, or
// [input0, input1, min0, min1, max0, max1] for separate.
// The separate order is needed because it's the only way to specify unknown
// numbers of inputs for ops like Concat.
enum { CONTIGUOUS_MIN_MAX, SEPARATE_MIN_MAX } min_max_order;
};
// Every op that has a quantized equivalent should be listed here, so that the
// conversion process can transform them.
const std::vector<QuantizedOpInfo>& GetQuantizedOpList() {
static const std::vector<QuantizedOpInfo> op_list = {
{"Add",
{},
{{"T1", DT_QUINT8}, {"T2", DT_QUINT8}, {"Toutput", DT_QINT32}},
DT_QUINT8,
DT_QINT32,
{},
QuantizedOpInfo::CONTIGUOUS_MIN_MAX},
{"AvgPool",
{"ksize", "strides", "padding"},
{{"T", DT_QUINT8}},
DT_QUINT8,
DT_QUINT8,
{},
QuantizedOpInfo::CONTIGUOUS_MIN_MAX},
{"BiasAdd",
{},
{{"T1", DT_QUINT8}, {"T2", DT_QUINT8}, {"out_type", DT_QINT32}},
DT_QUINT8,
DT_QINT32,
{},
QuantizedOpInfo::CONTIGUOUS_MIN_MAX},
{"Concat",
{"N"},
{{"T", DT_QUINT8}},
DT_QUINT8,
DT_QUINT8,
{0},
QuantizedOpInfo::SEPARATE_MIN_MAX},
{"Conv2D",
{"strides", "padding"},
{{"Tinput", DT_QUINT8}, {"Tfilter", DT_QUINT8}, {"out_type", DT_QINT32}},
DT_QUINT8,
DT_QINT32,
{},
QuantizedOpInfo::CONTIGUOUS_MIN_MAX},
{"MatMul",
{"transpose_a", "transpose_b"},
{{"T1", DT_QUINT8}, {"T2", DT_QUINT8}, {"Toutput", DT_QINT32}},
DT_QUINT8,
DT_QINT32,
{},
QuantizedOpInfo::CONTIGUOUS_MIN_MAX},
{"MaxPool",
{"ksize", "strides", "padding"},
{{"T", DT_QUINT8}},
DT_QUINT8,
DT_QUINT8,
{},
QuantizedOpInfo::CONTIGUOUS_MIN_MAX},
{"Mul",
{},
{{"T1", DT_QUINT8}, {"T2", DT_QUINT8}, {"Toutput", DT_QINT32}},
DT_QUINT8,
DT_QINT32,
{},
QuantizedOpInfo::CONTIGUOUS_MIN_MAX},
{"Relu",
{},
{{"Tinput", DT_QUINT8}},
DT_QUINT8,
DT_QUINT8,
{},
QuantizedOpInfo::CONTIGUOUS_MIN_MAX},
{"ResizeBilinear",
{"align_corners"},
{{"T", DT_QUINT8}},
DT_QUINT8,
DT_QUINT8,
{1},
QuantizedOpInfo::CONTIGUOUS_MIN_MAX},
{"Relu6",
{},
{{"Tinput", DT_QUINT8}},
DT_QUINT8,
DT_QUINT8,
{},
QuantizedOpInfo::CONTIGUOUS_MIN_MAX},
{"Reshape",
{},
{{"T", DT_QUINT8}},
DT_QUINT8,
DT_QUINT8,
{1},
QuantizedOpInfo::CONTIGUOUS_MIN_MAX},
};
return op_list;
}
namespace {
// Replaces invalid characters in input names to get a unique node name.
std::string UniqueNodeNameFromInput(const std::string& input_name) {
std::string prefix;
std::string node_name;
std::string suffix;
NodeNamePartsFromInput(input_name, &prefix, &node_name, &suffix);
std::string result;
if (prefix == "^") {
result += "__hat__";
}
result += node_name;
if (!suffix.empty()) {
result += "__port__" + suffix.substr(1, suffix.size() - 1);
}
return result;
}
// Pulls two float values from the named parameters, with a lot of checking.
absl::Status ExtractRangeFromParams(const TransformFuncContext& context,
const std::string& min_name,
const std::string& max_name,
float* min_value, float* max_value,
bool* has_range) {
// See if we've been given quantized inputs with a known range.
const bool has_min = (context.params.count(min_name) != 0);
const bool has_max = (context.params.count(max_name) != 0);
*has_range = (has_min || has_max);
if (!*has_range) {
return absl::OkStatus();
}
if (!has_min || !has_max) {
return absl::InvalidArgumentError(absl::StrCat("You must pass both ",
min_name, " and ", max_name,
" into quantize_nodes"));
}
TF_RETURN_IF_ERROR(context.GetOneFloatParameter(min_name, 0.0f, min_value));
TF_RETURN_IF_ERROR(context.GetOneFloatParameter(max_name, 0.0f, max_value));
return absl::OkStatus();
}
} // namespace
// Analyzes all the nodes in the graph to figure out which ones are duplicates
// apart from their names. This commonly includes identical Const nodes, but can
// also be simple operations that are repeated on multiple outputs of a
// particular node. The complexity is managed using a hash function that avoids
// the need for any O(n^2) algorithms when identifying duplicates.
absl::Status MergeDuplicateNodes(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def) {
// Make sure we can look up inputs and outputs quickly.
std::set<std::string> input_names(context.input_names.begin(),
context.input_names.end());
std::set<std::string> output_names(context.output_names.begin(),
context.output_names.end());
GraphDef current_graph_def = input_graph_def;
// Keep running the merging until no more duplicates are found.
bool any_duplicates_found;
do {
any_duplicates_found = false;
// First arrange all of the nodes by a hash of their contents.
std::map<uint64_t, std::vector<const NodeDef*>> hashed_nodes;
for (const NodeDef& node : current_graph_def.node()) {
NodeDef nameless_node = node;
// The name matters if it's being used as an input or output node,
// otherwise ignore it when looking for duplicates.
if (!input_names.count(node.name()) && !output_names.count(node.name())) {
nameless_node.set_name("");
}
const uint64_t hash = HashNodeDef(nameless_node);
hashed_nodes[hash].push_back(&node);
}
// If we have multiple nodes with the same hash, then we know they're
// duplicates and can be removed, unless they're stateful.
std::map<std::string, std::string> inputs_to_rename;
GraphDef merged_graph_def;
for (const std::pair<const uint64_t, std::vector<const NodeDef*>>&
hashed_node_info : hashed_nodes) {
const std::vector<const NodeDef*>& hash_node_list =
hashed_node_info.second;
for (int i = 0; i < hash_node_list.size(); ++i) {
const NodeDef* current_node = hash_node_list[i];
const OpDef* op_def = nullptr;
TF_RETURN_IF_ERROR(
OpRegistry::Global()->LookUpOpDef(current_node->op(), &op_def));
const bool is_duplicate = ((!op_def->is_stateful()) && (i > 0));
if (is_duplicate) {
const std::string original_name = hash_node_list[0]->name();
inputs_to_rename[current_node->name() + ":*"] = original_name;
any_duplicates_found = true;
} else {
NodeDef* new_node = merged_graph_def.mutable_node()->Add();
*new_node = *current_node;
}
}
}
// Update the graph so that any nodes that referred to removed inputs now
// pull from the remaining duplicate.
TF_RETURN_IF_ERROR(RenameNodeInputs(merged_graph_def, inputs_to_rename,
std::unordered_set<std::string>(),
&current_graph_def));
} while (any_duplicates_found);
*output_graph_def = current_graph_def;
return absl::OkStatus();
}
// Looks for the patterns that indicate there are two eight-bit ops feeding into
// each other, separated by a conversion up to float and back again. These occur
// during the initial conversion of ops to their quantized forms. Because we're
// only looking at an individual op in that phase and don't know if its inputs
// and outputs are eight-bit-capable, we start by converting the actual op into
// quantized form, but add float conversions before and after. This pass gets
// rid of those conversions if it turns out we do have adjacent ops capable of
// eight-bit processing.
absl::Status RemoveRedundantQuantizations(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def) {
std::set<std::string> graph_outputs;
for (const std::string& output_name : context.output_names) {
graph_outputs.insert(NodeNameFromInput(output_name));
}
std::map<std::string, std::string> inputs_to_rename;
GraphDef replaced_graph_def;
TF_RETURN_IF_ERROR(ReplaceMatchingOpTypes(
input_graph_def, // clang-format off
{"QuantizeV2",
{
{"Dequantize"},
{"Min"},
{"Max"},
}
}, // clang-format on
[&inputs_to_rename, &graph_outputs](
const NodeMatch& match, const std::set<std::string>& input_nodes,
const std::set<std::string>& output_nodes,
std::vector<NodeDef>* new_nodes) {
const NodeDef& quantize_node = match.node;
const NodeDef& dequantize_node = match.inputs[0].node;
inputs_to_rename[quantize_node.name() + ":0"] =
dequantize_node.input(0);
inputs_to_rename[quantize_node.name() + ":1"] =
dequantize_node.input(1);
inputs_to_rename[quantize_node.name() + ":2"] =
dequantize_node.input(2);
// Are other sub-graphs using the float intermediate result? If so,
// preserve it, but the input renaming still rewires the eight-bit ops
// so they don't go through float.
if (output_nodes.count(dequantize_node.name()) ||
graph_outputs.count(dequantize_node.name())) {
CopyOriginalMatch(match, new_nodes);
}
return absl::OkStatus();
},
{true}, &replaced_graph_def));
return RenameNodeInputs(replaced_graph_def, inputs_to_rename,
std::unordered_set<std::string>(), output_graph_def);
}
// If the user has passed in the input_min and input_max args, then we need to
// convert any input placeholders from float to eight bit, so quantized inputs
// can be fed directly into the graph.
absl::Status QuantizePlaceholders(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def) {
float input_min;
float input_max;
bool has_input_range;
TF_RETURN_IF_ERROR(ExtractRangeFromParams(context, "input_min", "input_max",
&input_min, &input_max,
&has_input_range));
if (!has_input_range) {
*output_graph_def = input_graph_def;
return absl::OkStatus();
}
std::map<std::string, std::string> inputs_to_rename_first_pass;
std::map<std::string, std::string> inputs_to_rename_second_pass;
GraphDef placeholder_graph_def;
placeholder_graph_def.Clear();
for (const NodeDef& node : input_graph_def.node()) {
if (node.op() != "Placeholder") {
*(placeholder_graph_def.mutable_node()->Add()) = node;
} else {
std::string namespace_prefix = node.name() + "_eightbit";
NodeDef quantized_placeholder;
quantized_placeholder = node;
SetNodeAttr("dtype", DT_QUINT8, &quantized_placeholder);
*(placeholder_graph_def.mutable_node()->Add()) = quantized_placeholder;
NodeDef min_node;
min_node.set_op("Const");
min_node.set_name(namespace_prefix + "/min");
SetNodeAttr("dtype", DT_FLOAT, &min_node);
Tensor min_tensor(DT_FLOAT, {});
min_tensor.flat<float>()(0) = input_min;
SetNodeTensorAttr<float>("value", min_tensor, &min_node);
*(placeholder_graph_def.mutable_node()->Add()) = min_node;
NodeDef max_node;
max_node.set_op("Const");
max_node.set_name(namespace_prefix + "/max");
SetNodeAttr("dtype", DT_FLOAT, &max_node);
Tensor max_tensor(DT_FLOAT, {});
max_tensor.flat<float>()(0) = input_max;
SetNodeTensorAttr<float>("value", max_tensor, &max_node);
*(placeholder_graph_def.mutable_node()->Add()) = max_node;
const std::string rename_suffix = "__RENAMED_PLACEHOLDER__";
NodeDef dequantize_node;
dequantize_node.set_op("Dequantize");
dequantize_node.set_name(namespace_prefix + "/dequantize");
SetNodeAttr("T", DT_QUINT8, &dequantize_node);
SetNodeAttr("mode", "MIN_FIRST", &dequantize_node);
AddNodeInput(node.name() + rename_suffix, &dequantize_node);
AddNodeInput(min_node.name(), &dequantize_node);
AddNodeInput(max_node.name(), &dequantize_node);
*(placeholder_graph_def.mutable_node()->Add()) = dequantize_node;
// First make sure that any internal references to the old placeholder
// now point to the dequantize result.
inputs_to_rename_first_pass[node.name()] = dequantize_node.name();
// Then fix up the dequantize op so that it really points to the
// placeholder.
inputs_to_rename_second_pass[node.name() + rename_suffix] = node.name();
}
}
GraphDef first_pass_graph_def;
TF_RETURN_IF_ERROR(RenameNodeInputs(
placeholder_graph_def, inputs_to_rename_first_pass,
std::unordered_set<std::string>(), &first_pass_graph_def));
TF_RETURN_IF_ERROR(
RenameNodeInputs(first_pass_graph_def, inputs_to_rename_second_pass,
std::unordered_set<std::string>(), output_graph_def));
return absl::OkStatus();
}
// During training, FakeQuantWithMinMaxVars ops capture a good min/max range for
// an activation layer. To use these during inference, this pass converts those
// ops into Requantizes with the trained min/maxes as constant inputs.
absl::Status ConvertFakeQuantsToRequantize(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def) {
TF_RETURN_IF_ERROR(ReplaceMatchingOpTypes(
input_graph_def, // clang-format off
{"FakeQuantWithMinMaxVars",
{
{"*"},
{"Const"},
{"Const"},
}
}, // clang-format on
[](const NodeMatch& match, const std::set<std::string>& input_nodes,
const std::set<std::string>& output_nodes,
std::vector<NodeDef>* new_nodes) {
const NodeDef& fake_quant_node = match.node;
const NodeDef& original_op_node = match.inputs[0].node;
const NodeDef& fake_quant_min_node = match.inputs[1].node;
const NodeDef& fake_quant_max_node = match.inputs[2].node;
std::string namespace_prefix = fake_quant_node.name() + "_eightbit";
new_nodes->push_back(original_op_node);
new_nodes->push_back(fake_quant_min_node);
new_nodes->push_back(fake_quant_max_node);
NodeDef quantize_node;
quantize_node.set_op("QuantizeV2");
quantize_node.set_name(namespace_prefix + "/quantize");
SetNodeAttr("T", DT_QINT32, &quantize_node);
SetNodeAttr("mode", "MIN_FIRST", &quantize_node);
AddNodeInput(fake_quant_node.input(0), &quantize_node);
AddNodeInput(fake_quant_min_node.name(), &quantize_node);
AddNodeInput(fake_quant_max_node.name(), &quantize_node);
new_nodes->push_back(quantize_node);
NodeDef requantize_node;
requantize_node.set_op("Requantize");
requantize_node.set_name(namespace_prefix + "/requantize");
SetNodeAttr("Tinput", DT_QINT32, &requantize_node);
SetNodeAttr("out_type", DT_QUINT8, &requantize_node);
AddNodeInput(quantize_node.name() + ":0", &requantize_node);
AddNodeInput(quantize_node.name() + ":1", &requantize_node);
AddNodeInput(quantize_node.name() + ":2", &requantize_node);
AddNodeInput(fake_quant_min_node.name(), &requantize_node);
AddNodeInput(fake_quant_max_node.name(), &requantize_node);
new_nodes->push_back(requantize_node);
// Convert the 8-bit result back into float for the final output.
NodeDef dequantize_node;
dequantize_node.set_op("Dequantize");
dequantize_node.set_name(fake_quant_node.name());
SetNodeAttr("T", DT_QUINT8, &dequantize_node);
SetNodeAttr("mode", "MIN_FIRST", &dequantize_node);
AddNodeInput(requantize_node.name() + ":0", &dequantize_node);
AddNodeInput(requantize_node.name() + ":1", &dequantize_node);
AddNodeInput(requantize_node.name() + ":2", &dequantize_node);
new_nodes->push_back(dequantize_node);
return absl::OkStatus();
},
{}, output_graph_def));
return absl::OkStatus();
}
// We always generate Requantize ops driven by dynamic RequantizationRange
// calculations when we produce quantized ops like Conv2D or BiasAdd with
// 32-bit results. If there were FakeQuant ops already for those activation
// layers, then there will be a later Requantize op with constant min/max
// inputs, which is preferable for fast inference. This pass looks for those
// later Requantize ops, and replaces the dynamic version with them.
absl::Status MergeAdjacentRequantizes(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def) {
TF_RETURN_IF_ERROR(ReplaceMatchingOpTypes(
input_graph_def, // clang-format off
{"Requantize",
{
{"QuantizeV2",
{
{"Dequantize",
{
{"Requantize",
{
{"*"},
{"*"},
{"*"},
{"RequantizationRange"},
{"RequantizationRange"},
}
},
{"Requantize"},
{"Requantize"},
}
},
{"Const"},
{"Const"},
},
},
{"QuantizeV2"},
{"QuantizeV2"},
{"Const"},
{"Const"},
}
}, // clang-format on
[](const NodeMatch& match, const std::set<std::string>& input_nodes,
const std::set<std::string>& output_nodes,
std::vector<NodeDef>* new_nodes) {
const NodeDef& fake_requantize_node = match.node;
const NodeDef& original_op_node =
match.inputs[0].inputs[0].inputs[0].inputs[0].node;
const NodeDef& fake_requantize_min_node = match.inputs[3].node;
const NodeDef& fake_requantize_max_node = match.inputs[4].node;
new_nodes->push_back(original_op_node);
new_nodes->push_back(fake_requantize_min_node);
new_nodes->push_back(fake_requantize_max_node);
NodeDef requantize_node;
requantize_node = fake_requantize_node;
requantize_node.mutable_input()->Clear();
AddNodeInput(original_op_node.name() + ":0", &requantize_node);
AddNodeInput(original_op_node.name() + ":1", &requantize_node);
AddNodeInput(original_op_node.name() + ":2", &requantize_node);
AddNodeInput(fake_requantize_min_node.name(), &requantize_node);
AddNodeInput(fake_requantize_max_node.name(), &requantize_node);
new_nodes->push_back(requantize_node);
return absl::OkStatus();
},
{}, output_graph_def));
return absl::OkStatus();
}
// Sometimes FakeQuantWithMinMaxVars ops are added at the end of a chain of
// linear ops like Relu, MaxPool, etc, several steps from the Conv2D or BiasAdd
// op that we want to apply the trained constant conversions to. This pass tries
// to move FakeQuant ops up the input chain, so they're as close as possible to
// the 32-bit conversion, and so can be easily merged into the automatic dynamic
// Requantizes.
absl::Status HoistFakeQuants(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def) {
GraphDef current_graph_def = input_graph_def;
const int max_depth = 3;
for (int depth = max_depth; depth > 0; --depth) {
OpTypePattern pattern = {"*"};
for (int i = 0; i < depth; ++i) {
pattern = {"*", {pattern}};
}
pattern = {"FakeQuantWithMinMaxVars", {pattern, {"Const"}, {"Const"}}};
GraphDef hoisted_graph_def;
TF_RETURN_IF_ERROR(ReplaceMatchingOpTypes(
current_graph_def, pattern,
[depth](const NodeMatch& match,
const std::set<std::string>& input_nodes,
const std::set<std::string>& output_nodes,
std::vector<NodeDef>* new_nodes) {
const NodeDef& fake_quant_node = match.node;
const NodeDef& fake_quant_min_node = match.inputs[1].node;
const NodeDef& fake_quant_max_node = match.inputs[2].node;
std::vector<NodeDef> linear_nodes;
NodeMatch current_match = match;
for (int i = 0; i <= depth; ++i) {
linear_nodes.push_back(current_match.inputs[0].node);
current_match = current_match.inputs[0];
}
NodeDef new_fake_quant_node;
new_fake_quant_node = fake_quant_node;
new_fake_quant_node.set_name(fake_quant_node.name() + "_hoisted");
new_fake_quant_node.set_input(
0, linear_nodes[linear_nodes.size() - 2].input(0));
new_nodes->push_back(new_fake_quant_node);
new_nodes->push_back(fake_quant_min_node);
new_nodes->push_back(fake_quant_max_node);
linear_nodes[linear_nodes.size() - 2].set_input(
0, new_fake_quant_node.name());
linear_nodes.front().set_name(fake_quant_node.name());
for (const NodeDef& linear_node : linear_nodes) {
new_nodes->push_back(linear_node);
}
return absl::OkStatus();
},
{}, &hoisted_graph_def));
current_graph_def = hoisted_graph_def;
}
*output_graph_def = current_graph_def;
return absl::OkStatus();
}
// Converts any float ops that have eight-bit equivalents into their quantized
// forms, so that as much calculation as possible is done in the lower-precision
// format.
absl::Status QuantizeNodes(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def) {
// Loop through all of the quantizable op types, and replace any occurrences
// with equivalent sub-graphs with quantized ops at their core. For example
// this one-input operation:
//
// Input(float)
// |
// v
// Operation
// |
// v
// (float)
//
// Will be turned into it's quantized equivalent:
//
// Input(float) ReshapeDims
// +------v v-------------+
// | Reshape
// | |
// | | ReductionDims
// | +-----+ |
// | | +---c---------+
// | v v v v-------+
// | Min Max
// | +----+ |
// v v v--------+
// Quantize
// |
// v
// QuantizedOperation
// | | |
// v v v
// Dequantize
// |
// v
// (float)
//
// This keeps the inputs and outputs visible to the rest of the graph in
// float
// and converts them down to quantized buffers internally for the
// computation.
// The result will end up with a lot of redundant dequantize/quantize pairs
// between adjacent quantized ops, but a later pass removes these where it
// can.
std::set<std::string> ops_to_ignore;
if (context.params.count("ignore_op") > 0) {
for (const std::string& name : context.params.at("ignore_op")) {
ops_to_ignore.insert(name);
}
}
const std::vector<QuantizedOpInfo>& op_list = GetQuantizedOpList();
std::string op_pattern;
bool is_first = true;
std::map<std::string, QuantizedOpInfo> op_map;
for (const QuantizedOpInfo& op_info : op_list) {
if (ops_to_ignore.count(op_info.float_name) == 0) {
absl::StrAppend(&op_pattern, is_first ? "" : "|", op_info.float_name);
op_map.insert({op_info.float_name, op_info});
is_first = false;
}
}
// If input_min and input max have been passed in, then we convert all float
// Placeholder nodes into quantized versions, with the supplied values as
// their range.
GraphDef placeholder_graph_def;
TF_RETURN_IF_ERROR(
QuantizePlaceholders(input_graph_def, context, &placeholder_graph_def));
TF_RETURN_IF_ERROR(IsGraphValid(placeholder_graph_def));
// If there are any FakeQuantWithMinMaxVars at the end of a chain of linear
// operations like Relu or MaxPool, move them up so that they're as close as
// possible to ops with 32-bit outputs like BiasAdd or Conv2D.
GraphDef hoisted_graph_def;
TF_RETURN_IF_ERROR(
HoistFakeQuants(placeholder_graph_def, context, &hoisted_graph_def));
TF_RETURN_IF_ERROR(IsGraphValid(hoisted_graph_def));
// Convert any FakeQuantWithMinMaxVars, which hold the trained ranges of
// activation layers, into Requantize ops with those ranges instead. This
// makes it easier to replace the dynamic range calculations that are used
// by default.
GraphDef converted_graph_def;
TF_RETURN_IF_ERROR(ConvertFakeQuantsToRequantize(hoisted_graph_def, context,
&converted_graph_def));
TF_RETURN_IF_ERROR(IsGraphValid(converted_graph_def));
// If fallback_min and fallback_max are set, then we'll use hardwired ranges
// for all the 32-bit to 8-bit requantizations.
float fallback_min;
float fallback_max;
bool has_fallback_range;
TF_RETURN_IF_ERROR(ExtractRangeFromParams(
context, "fallback_min", "fallback_max", &fallback_min, &fallback_max,
&has_fallback_range));
// Replace all occurrences of the current float op with its quantized
// equivalent.
GraphDef quantized_graph_def;
TF_RETURN_IF_ERROR(ReplaceMatchingOpTypes(
converted_graph_def, {op_pattern},
[&op_map, fallback_min, fallback_max, has_fallback_range](
const NodeMatch& match, const std::set<std::string>& input_nodes,
const std::set<std::string>& output_nodes,
std::vector<NodeDef>* new_nodes) {
const NodeDef& float_node = match.node;
const QuantizedOpInfo& op_info = op_map[float_node.op()];
DataTypeVector input_types;
DataTypeVector output_types;
TF_RETURN_IF_ERROR(
GetInOutTypes(float_node, &input_types, &output_types));
bool are_all_float = true;
for (int i = 0; i < float_node.input_size(); ++i) {
// Skip any known non-float inputs.
if (op_info.unquantized_inputs.count(i)) {
continue;
}
if (i >= input_types.size()) {
LOG(ERROR) << "input_types has incorrect size "
<< input_types.size() << " <= " << i
<< ". Assuming everything else is floats.";
}
if (i < input_types.size() && input_types[i] != DT_FLOAT) {
are_all_float = false;
}
}
for (const DataType& output_type : output_types) {
if (output_type != DT_FLOAT) {
are_all_float = false;
}
}
// This isn't a float op, so don't quantize it.
if (!are_all_float) {
CopyOriginalMatch(match, new_nodes);
return absl::OkStatus();
}
std::string namespace_prefix = float_node.name() + "_eightbit";
// Quantize all of the inputs.
std::vector<std::string> quantized_input_names;
for (int i = 0; i < float_node.input_size(); ++i) {
// Skip any non-float inputs.
if (op_info.unquantized_inputs.count(i)) {
continue;
}
const std::string& input_name = float_node.input(i);
std::string unique_input_name =
namespace_prefix + "/" + UniqueNodeNameFromInput(input_name);
// Add some common constants we need for reshaping inputs.
NodeDef reshape_dims;
reshape_dims.set_op("Const");
reshape_dims.set_name(unique_input_name + "/reshape_dims");
AddNodeInput("^" + NodeNameFromInput(input_name), &reshape_dims);
SetNodeAttr("dtype", DT_INT32, &reshape_dims);
Tensor reshape_dims_tensor(DT_INT32, {1});
reshape_dims_tensor.flat<int32_t>()(0) = -1;
SetNodeTensorAttr<int32_t>("value", reshape_dims_tensor,
&reshape_dims);
new_nodes->push_back(reshape_dims);
NodeDef reduction_dims;
reduction_dims.set_op("Const");
reduction_dims.set_name(unique_input_name + "/reduction_dims");
AddNodeInput("^" + NodeNameFromInput(input_name), &reduction_dims);
SetNodeAttr("dtype", DT_INT32, &reduction_dims);
Tensor reduction_dims_tensor(DT_INT32, {1});
reduction_dims_tensor.flat<int32_t>()(0) = 0;
SetNodeTensorAttr<int32_t>("value", reduction_dims_tensor,
&reduction_dims);
new_nodes->push_back(reduction_dims);
NodeDef reshape_node;
reshape_node.set_op("Reshape");
reshape_node.set_name(unique_input_name + "/reshape");
SetNodeAttr("T", DT_FLOAT, &reshape_node);
AddNodeInput(input_name, &reshape_node);
AddNodeInput(reshape_dims.name(), &reshape_node);
new_nodes->push_back(reshape_node);
NodeDef min_node;
min_node.set_op("Min");
min_node.set_name(unique_input_name + "/min");
SetNodeAttr("T", DT_FLOAT, &min_node);
SetNodeAttr("keep_dims", false, &min_node);
AddNodeInput(reshape_node.name(), &min_node);
AddNodeInput(reduction_dims.name(), &min_node);
new_nodes->push_back(min_node);
NodeDef max_node;
max_node.set_op("Max");
max_node.set_name(unique_input_name + "/max");
SetNodeAttr("T", DT_FLOAT, &max_node);
SetNodeAttr("keep_dims", false, &max_node);
AddNodeInput(reshape_node.name(), &max_node);
AddNodeInput(reduction_dims.name(), &max_node);
new_nodes->push_back(max_node);
NodeDef quantize_node;
quantize_node.set_op("QuantizeV2");
quantize_node.set_name(unique_input_name + "/quantize");
SetNodeAttr("T", DT_QUINT8, &quantize_node);
SetNodeAttr("mode", "MIN_FIRST", &quantize_node);
AddNodeInput(input_name, &quantize_node);
AddNodeInput(min_node.name(), &quantize_node);
AddNodeInput(max_node.name(), &quantize_node);
new_nodes->push_back(quantize_node);
quantized_input_names.push_back(quantize_node.name());
}
// Set up the quantized version of the current op.
NodeDef quantized_main_node;
quantized_main_node.set_op("Quantized" + float_node.op());
quantized_main_node.set_name(float_node.name() + "/eightbit");
for (const std::string& attr_to_copy : op_info.attrs_to_copy) {
CopyNodeAttr(float_node, attr_to_copy, attr_to_copy,
&quantized_main_node);
}
for (const std::pair<std::string, DataType>& dtype_to_set :
op_info.dtypes_to_set) {
SetNodeAttr(dtype_to_set.first, dtype_to_set.second,
&quantized_main_node);
}
int quantized_input_index = 0;
for (int i = 0; i < float_node.input_size(); ++i) {
if (op_info.unquantized_inputs.count(i)) {
AddNodeInput(float_node.input(i), &quantized_main_node);
} else {
const std::string& quantized_input_name =
quantized_input_names[quantized_input_index];
AddNodeInput(quantized_input_name + ":0", &quantized_main_node);
++quantized_input_index;
}
}
if (op_info.min_max_order == QuantizedOpInfo::CONTIGUOUS_MIN_MAX) {
for (const std::string& quantized_input_name :
quantized_input_names) {
AddNodeInput(quantized_input_name + ":1", &quantized_main_node);
AddNodeInput(quantized_input_name + ":2", &quantized_main_node);
}
} else {
for (const std::string& quantized_input_name :
quantized_input_names) {
AddNodeInput(quantized_input_name + ":1", &quantized_main_node);
}
for (const std::string& quantized_input_name :
quantized_input_names) {
AddNodeInput(quantized_input_name + ":2", &quantized_main_node);
}
}
new_nodes->push_back(quantized_main_node);
std::string eight_bit_node_name;
if (op_info.output_bit_depth == DT_QINT32) {
// Shrink the range of the output down from 32 bits to 8.
std::string requantize_min_input;
std::string requantize_max_input;
if (has_fallback_range) {
// Use constant values for the min/max range if they were given.
NodeDef fallback_min_node;
fallback_min_node.set_op("Const");
fallback_min_node.set_name(quantized_main_node.name() +
"/fallback_min");
SetNodeAttr("dtype", DT_FLOAT, &fallback_min_node);
Tensor fallback_min_tensor(DT_FLOAT, {});
fallback_min_tensor.flat<float>()(0) = fallback_min;
SetNodeTensorAttr<float>("value", fallback_min_tensor,
&fallback_min_node);
new_nodes->push_back(fallback_min_node);
NodeDef fallback_max_node;
fallback_max_node.set_op("Const");
fallback_max_node.set_name(quantized_main_node.name() +
"/fallback_max");
SetNodeAttr("dtype", DT_FLOAT, &fallback_max_node);
Tensor fallback_max_tensor(DT_FLOAT, {});
fallback_max_tensor.flat<float>()(0) = fallback_max;
SetNodeTensorAttr<float>("value", fallback_max_tensor,
&fallback_max_node);
new_nodes->push_back(fallback_max_node);
requantize_min_input = fallback_min_node.name();
requantize_max_input = fallback_max_node.name();
} else {
// Otherwise dynamically measure the range each time.
NodeDef requant_range_node;
requant_range_node.set_op("RequantizationRange");
requant_range_node.set_name(quantized_main_node.name() +
"/requant_range");
SetNodeAttr("Tinput", DT_QINT32, &requant_range_node);
AddNodeInput(quantized_main_node.name() + ":0",
&requant_range_node);
AddNodeInput(quantized_main_node.name() + ":1",
&requant_range_node);
AddNodeInput(quantized_main_node.name() + ":2",
&requant_range_node);
new_nodes->push_back(requant_range_node);
requantize_min_input = requant_range_node.name() + ":0";
requantize_max_input = requant_range_node.name() + ":1";
}
NodeDef requantize_node;
requantize_node.set_op("Requantize");
requantize_node.set_name(quantized_main_node.name() + "/requantize");
SetNodeAttr("Tinput", DT_QINT32, &requantize_node);
SetNodeAttr("out_type", DT_QUINT8, &requantize_node);
AddNodeInput(quantized_main_node.name() + ":0", &requantize_node);
AddNodeInput(quantized_main_node.name() + ":1", &requantize_node);
AddNodeInput(quantized_main_node.name() + ":2", &requantize_node);
AddNodeInput(requantize_min_input, &requantize_node);
AddNodeInput(requantize_max_input, &requantize_node);
new_nodes->push_back(requantize_node);
eight_bit_node_name = requantize_node.name();
} else {
eight_bit_node_name = quantized_main_node.name();
}
// Convert the 8-bit result back into float for the final output.
NodeDef dequantize_node;
dequantize_node.set_op("Dequantize");
dequantize_node.set_name(float_node.name());
SetNodeAttr("T", DT_QUINT8, &dequantize_node);
SetNodeAttr("mode", "MIN_FIRST", &dequantize_node);
AddNodeInput(eight_bit_node_name + ":0", &dequantize_node);
AddNodeInput(eight_bit_node_name + ":1", &dequantize_node);
AddNodeInput(eight_bit_node_name + ":2", &dequantize_node);
new_nodes->push_back(dequantize_node);
return absl::OkStatus();
},
{}, &quantized_graph_def));
TF_RETURN_IF_ERROR(IsGraphValid(quantized_graph_def));
// If we've ended up with two Requantize ops in a row (for example if there
// was a Conv2D feeding into a FakeQuantWithMinMaxVars) merge them together,
// using the trained range from the second op.
GraphDef merged_graph_def;
TF_RETURN_IF_ERROR(MergeAdjacentRequantizes(quantized_graph_def, context,
&merged_graph_def));
TF_RETURN_IF_ERROR(IsGraphValid(merged_graph_def));
// There can be duplicate quantize nodes if multiple ops pull from a single
// input, which makes it harder to remove redundant ones, so strip them out.
GraphDef deduped_graph_def;
TF_RETURN_IF_ERROR(
MergeDuplicateNodes(merged_graph_def, context, &deduped_graph_def));
TF_RETURN_IF_ERROR(IsGraphValid(deduped_graph_def));
// Look for Dequantizes that immediately go into Quantizes, and remove them
// since the two together cancel each other out. This allows us to keep the
// data flow in eight bit where two adjacent ops are in eight bit, but still
// keep interoperability with float ops.
TF_RETURN_IF_ERROR(RemoveRedundantQuantizations(deduped_graph_def, context,
output_graph_def));
TF_RETURN_IF_ERROR(IsGraphValid(*output_graph_def));
return absl::OkStatus();
}
REGISTER_GRAPH_TRANSFORM("quantize_nodes", QuantizeNodes);
REGISTER_GRAPH_TRANSFORM("merge_duplicate_nodes", MergeDuplicateNodes);
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/quantize_nodes_test.cc
+1534
View File
File diff suppressed because it is too large Load Diff
tensorflow/tools/graph_transforms/quantize_weights.cc
+146
View File
@@ -0,0 +1,146 @@
/* Copyright 2016 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.
==============================================================================*/
#define EIGEN_USE_THREADS
#include "tensorflow/core/common_runtime/constant_folding.h"
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/common_runtime/threadpool_device.h"
#include "tensorflow/core/graph/node_builder.h"
#include "tensorflow/core/graph/subgraph.h"
#include "tensorflow/core/kernels/quantization_utils.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
// Converts any large float constants into eight-bit equivalents, with a
// Dequantize op so that subsequent nodes can still access the results in a
// float form.
absl::Status QuantizeWeights(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def) {
int32_t minimum_size;
TF_RETURN_IF_ERROR(
context.GetOneInt32Parameter("minimum_size", 1024, &minimum_size));
TF_RETURN_IF_ERROR(ReplaceMatchingOpTypes(
input_graph_def, {"Const"},
[minimum_size](const NodeMatch& match,
const std::set<std::string>& input_nodes,
const std::set<std::string>& output_nodes,
std::vector<NodeDef>* new_nodes) {
const NodeDef& old_const_node = match.node;
if (!old_const_node.attr().count("dtype")) {
return absl::InvalidArgumentError(absl::StrCat(
"No 'dtype' attribute for Const node ", old_const_node.name()));
}
if (!old_const_node.attr().count("value")) {
return absl::InvalidArgumentError(absl::StrCat(
"No 'value' attribute for Const node ", old_const_node.name()));
}
const DataType old_dtype = old_const_node.attr().at("dtype").type();
Tensor old_tensor;
if (!old_tensor.FromProto(old_const_node.attr().at("value").tensor())) {
return absl::InvalidArgumentError(absl::StrCat(
"Decoding Tensor failed for node", old_const_node.name()));
}
const size_t num_elements = old_tensor.NumElements();
// If this isn't a float constant, or it's too small, then reuse the
// same node with no changes.
if ((old_dtype != DT_FLOAT) || (num_elements < minimum_size)) {
new_nodes->push_back(old_const_node);
return absl::OkStatus();
}
const float* old_values = old_tensor.flat<float>().data();
float min = std::numeric_limits<float>::max();
float max = std::numeric_limits<float>::min();
for (int i = 0; i < num_elements; ++i) {
const float value = old_values[i];
min = std::min(min, value);
max = std::max(max, value);
}
// Make sure the quantization range includes 0.0f. Not all quantized
// Ops behave properly if 0.0f is not in the range.
min = std::min(min, 0.0f);
max = std::max(0.0f, max);
// min_value == max_value is a tricky case. It can occur for general
// tensors, and of course for scalars. The quantized ops cannot deal
// with this case, so we set max_value to something else.
// It's a tricky question what is the numerically best solution to
// deal with this degeneracy.
// TODO(petewarden): Better use a tolerance than a hard comparison?
if (min == max) {
if (std::abs(min) < 0.000001f) {
max = min + 1.0f;
} else if (min > 0) {
max = 2.0f * min;
} else {
max = min / 2.0f;
}
}
Tensor quantized_tensor(DT_QUINT8, old_tensor.shape());
FloatTensorToQuantizedInPlace<quint8>(old_tensor, min, max,
&quantized_tensor);
NodeDef quantized_const_node;
quantized_const_node.set_op("Const");
quantized_const_node.set_name(old_const_node.name() +
"_quantized_const");
SetNodeAttr("dtype", DT_QUINT8, &quantized_const_node);
SetNodeTensorAttr<float>("value", quantized_tensor,
&quantized_const_node);
new_nodes->push_back(quantized_const_node);
NodeDef min_node;
min_node.set_op("Const");
min_node.set_name(old_const_node.name() + "_quantized_min");
SetNodeAttr("dtype", DT_FLOAT, &min_node);
Tensor min_tensor(DT_FLOAT, {});
min_tensor.scalar<float>()() = min;
SetNodeTensorAttr<float>("value", min_tensor, &min_node);
new_nodes->push_back(min_node);
NodeDef max_node;
max_node.set_op("Const");
max_node.set_name(old_const_node.name() + "_quantized_max");
SetNodeAttr("dtype", DT_FLOAT, &max_node);
Tensor max_tensor(DT_FLOAT, {});
max_tensor.scalar<float>()() = max;
SetNodeTensorAttr<float>("value", max_tensor, &max_node);
new_nodes->push_back(max_node);
NodeDef dequantize_node;
dequantize_node.set_op("Dequantize");
dequantize_node.set_name(old_const_node.name());
SetNodeAttr("T", DT_QUINT8, &dequantize_node);
SetNodeAttr("mode", "MIN_FIRST", &dequantize_node);
AddNodeInput(quantized_const_node.name(), &dequantize_node);
AddNodeInput(min_node.name(), &dequantize_node);
AddNodeInput(max_node.name(), &dequantize_node);
new_nodes->push_back(dequantize_node);
return absl::OkStatus();
},
{}, output_graph_def));
return absl::OkStatus();
}
REGISTER_GRAPH_TRANSFORM("quantize_weights", QuantizeWeights);
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/quantize_weights_test.cc
+163
View File
@@ -0,0 +1,163 @@
/* Copyright 2015 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/cc/ops/const_op.h"
#include "tensorflow/cc/ops/image_ops.h"
#include "tensorflow/cc/ops/nn_ops.h"
#include "tensorflow/cc/ops/sendrecv_ops.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
// Declare here, so we don't need a public header.
absl::Status QuantizeWeights(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def);
class QuantizeWeightsTest : public ::testing::Test {
protected:
void BuildGraphDef(const TensorShape& input_shape,
std::initializer_list<float> input_values,
const TensorShape& weight_shape,
std::initializer_list<float> weight_values,
GraphDef* original_graph_def) {
auto root = tensorflow::Scope::DisabledShapeInferenceScope();
Tensor input_data(DT_FLOAT, input_shape);
test::FillValues<float>(&input_data, input_values);
Output input_op =
ops::Const(root.WithOpName("input_op"), Input::Initializer(input_data));
Tensor weights_data(DT_FLOAT, weight_shape);
test::FillValues<float>(&weights_data, weight_values);
Output weights_op = ops::Const(root.WithOpName("weights_op"),
Input::Initializer(weights_data));
Output conv_op = ops::Conv2D(root.WithOpName("output"), input_op,
weights_op, {1, 1, 1, 1}, "VALID");
TF_ASSERT_OK(root.ToGraphDef(original_graph_def));
}
void TestQuantizeWeights() {
GraphDef original_graph_def;
BuildGraphDef({1, 1, 6, 2},
{1.0f, 4.0f, 2.0f, 5.0f, 3.0f, 6.0f, -1.0f, -4.0f, -2.0f,
-5.0f, -3.0f, -6.0f},
{1, 2, 2, 10},
{1.0f, 2.0f, 3.0f, 4.0f, 0.1f, 0.2f, 0.3f, 0.4f, 1.0f, 2.0f,
3.0f, 4.0f, 0.1f, 0.2f, 0.3f, 0.4f, 1.0f, 2.0f, 3.0f, 4.0f,
0.1f, 0.2f, 0.3f, 0.4f, 1.0f, 2.0f, 3.0f, 4.0f, 0.1f, 0.2f,
0.3f, 0.4f, 1.0f, 2.0f, 3.0f, 4.0f, 0.1f, 0.2f, 0.3f, 0.4f},
&original_graph_def);
TransformFuncContext context;
context.output_names = {"output"};
context.params["minimum_size"] = {"16"};
GraphDef quantized_graph_def;
TF_ASSERT_OK(
QuantizeWeights(original_graph_def, context, &quantized_graph_def));
// Verify the structure of the quantized graph.
std::map<std::string, const NodeDef*> node_lookup;
MapNamesToNodes(quantized_graph_def, &node_lookup);
EXPECT_EQ(1, node_lookup.count("input_op"));
const NodeDef* q_input_op = node_lookup.at("input_op");
EXPECT_EQ(DT_FLOAT, q_input_op->attr().at("dtype").type());
EXPECT_EQ(1, node_lookup.count("weights_op"));
const NodeDef* q_weights_op = node_lookup.at("weights_op");
EXPECT_EQ("Dequantize", q_weights_op->op());
const std::string& weights_const_name =
NodeNameFromInput(q_weights_op->input(0));
EXPECT_EQ(1, node_lookup.count(weights_const_name));
const NodeDef* q_weights_const = node_lookup.at(weights_const_name);
EXPECT_EQ("Const", q_weights_const->op());
EXPECT_EQ(DT_QUINT8, q_weights_const->attr().at("dtype").type());
// Run the original graph.
std::unique_ptr<Session> original_session(NewSession(SessionOptions()));
TF_ASSERT_OK(original_session->Create(original_graph_def));
std::vector<Tensor> original_outputs;
TF_ASSERT_OK(original_session->Run({}, {"output"}, {}, &original_outputs));
// Run the quantized graph.
std::unique_ptr<Session> quantized_session(NewSession(SessionOptions()));
TF_ASSERT_OK(quantized_session->Create(quantized_graph_def));
std::vector<Tensor> quantized_outputs;
TF_ASSERT_OK(
quantized_session->Run({}, {"output"}, {}, &quantized_outputs));
// Compare the results
test::ExpectTensorNear<float>(original_outputs[0], quantized_outputs[0],
0.5);
}
};
TEST_F(QuantizeWeightsTest, TestQuantizeWeights) { TestQuantizeWeights(); }
TEST_F(QuantizeWeightsTest, RangesAlwaysIncludeZero) {
GraphDef original_graph_def;
BuildGraphDef({1, 1, 4, 4},
{-1.0f, -4.0f, -2.0f, -5.0f, -1.0f, -4.0f, -2.0f, -5.0f, -1.0f,
-4.0f, -2.0f, -5.0f, -1.0f, -4.0f, -2.0f, -5.0f},
{1, 2, 2, 10},
{1.0f, 2.0f, 3.0f, 4.0f, 0.1f, 0.2f, 0.3f, 0.4f, 1.0f, 2.0f,
3.0f, 4.0f, 0.1f, 0.2f, 0.3f, 0.4f, 1.0f, 2.0f, 3.0f, 4.0f,
0.1f, 0.2f, 0.3f, 0.4f, 1.0f, 2.0f, 3.0f, 4.0f, 0.1f, 0.2f,
0.3f, 0.4f, 1.0f, 2.0f, 3.0f, 4.0f, 0.1f, 0.2f, 0.3f, 0.4f},
&original_graph_def);
TransformFuncContext context;
context.output_names = {"output"};
context.params["minimum_size"] = {"16"};
GraphDef quantized_graph_def;
TF_ASSERT_OK(
QuantizeWeights(original_graph_def, context, &quantized_graph_def));
std::map<std::string, const NodeDef*> node_lookup;
MapNamesToNodes(quantized_graph_def, &node_lookup);
auto expected_tensor = [](float value) {
Tensor tensor(DT_FLOAT, TensorShape({}));
test::FillValues<float>(&tensor, {value});
return tensor;
};
auto existing_tensor = [&node_lookup](std::string op) {
const NodeDef* node_def = node_lookup.at(op);
CHECK(node_def);
return GetNodeTensorAttr(*node_def, "value");
};
// The max of input_op is moved from -1.0 to 0.0.
test::ExpectTensorNear<float>(
expected_tensor(-5.0), existing_tensor("input_op_quantized_min"), 1e-5);
test::ExpectTensorNear<float>(
expected_tensor(0.0), existing_tensor("input_op_quantized_max"), 1e-5);
// The min of weights_op is moved from 0.1 to 0.0.
test::ExpectTensorNear<float>(
expected_tensor(0.0), existing_tensor("weights_op_quantized_min"), 1e-5);
test::ExpectTensorNear<float>(
expected_tensor(4.0), existing_tensor("weights_op_quantized_max"), 1e-5);
}
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/remove_attribute.cc
+68
View File
@@ -0,0 +1,68 @@
/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/core/common_runtime/constant_folding.h"
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/graph/node_builder.h"
#include "tensorflow/core/graph/subgraph.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/tools/graph_transforms/fold_constants_lib.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
// Deletes a given attribute from the specified nodes.
absl::Status RemoveAttribute(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def) {
if (!context.params.count("attribute_name") ||
(context.params.at("attribute_name").size() != 1)) {
return absl::InvalidArgumentError(
"remove_attribute expects exactly one 'attribute_name' "
"argument, e.g. remove_attribute(op_name=Mul, attribute_name=foo)");
}
std::string op_name;
if (context.params.count("op_name")) {
if (context.params.at("op_name").size() != 1) {
return absl::InvalidArgumentError(absl::StrCat(
"remove_attribute expects a single op_name argument, but found ",
context.params.at("op_name").size()));
}
op_name = context.params.at("op_name")[0];
} else {
op_name = "*";
}
const std::string attribute_name = context.params.at("attribute_name")[0];
output_graph_def->Clear();
for (const NodeDef& node : input_graph_def.node()) {
NodeDef* new_node = output_graph_def->mutable_node()->Add();
*new_node = node;
if (((op_name == "*") || (op_name == node.op())) &&
(node.attr().count(attribute_name))) {
new_node->mutable_attr()->erase(attribute_name);
}
}
return absl::OkStatus();
}
REGISTER_GRAPH_TRANSFORM("remove_attribute", RemoveAttribute);
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/remove_attribute_test.cc
+125
View File
@@ -0,0 +1,125 @@
/* Copyright 2015 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/cc/ops/const_op.h"
#include "tensorflow/cc/ops/image_ops.h"
#include "tensorflow/cc/ops/nn_ops.h"
#include "tensorflow/cc/ops/sendrecv_ops.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
// Declare here, so we don't need a public header.
absl::Status RemoveAttribute(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def);
class RemoveAttributeTest : public ::testing::Test {
protected:
void TestRemoveAttribute() {
GraphDef graph_def;
NodeDef* mul_node1 = graph_def.add_node();
mul_node1->set_name("mul_node1");
mul_node1->set_op("Mul");
mul_node1->add_input("add_node2");
mul_node1->add_input("add_node3");
SetNodeAttr<int32_t>("foo", 23, mul_node1);
SetNodeAttr<std::string>("bar", "something", mul_node1);
NodeDef* add_node2 = graph_def.add_node();
add_node2->set_name("add_node2");
add_node2->set_op("Add");
add_node2->add_input("const_node1");
add_node2->add_input("const_node2");
SetNodeAttr<int32_t>("foo", 46, add_node2);
SetNodeAttr<int32_t>("bob", 23, add_node2);
SetNodeAttr<std::string>("bar", "something else", add_node2);
NodeDef* add_node3 = graph_def.add_node();
add_node3->set_name("add_node3");
add_node3->set_op("Add");
add_node3->add_input("const_node1");
add_node3->add_input("const_node3");
NodeDef* const_node1 = graph_def.add_node();
const_node1->set_name("const_node1");
const_node1->set_op("Const");
NodeDef* const_node2 = graph_def.add_node();
const_node2->set_name("const_node2");
const_node2->set_op("Const");
NodeDef* const_node3 = graph_def.add_node();
const_node3->set_name("const_node3");
const_node3->set_op("Const");
NodeDef* add_node4 = graph_def.add_node();
add_node4->set_name("add_node4");
add_node4->set_op("Add");
add_node4->add_input("add_node2");
add_node4->add_input("add_node3");
GraphDef wildcard_result;
TransformFuncContext context;
context.input_names = {};
context.output_names = {"mul_node1"};
context.params.insert(std::pair<std::string, std::vector<std::string>>(
{"op_name", {std::string("*")}}));
context.params.insert(std::pair<std::string, std::vector<std::string>>(
{"attribute_name", {std::string("foo")}}));
TF_ASSERT_OK(RemoveAttribute(graph_def, context, &wildcard_result));
std::map<std::string, const NodeDef*> node_lookup;
MapNamesToNodes(wildcard_result, &node_lookup);
EXPECT_EQ(0, node_lookup.at("mul_node1")->attr().count("foo"));
EXPECT_EQ(1, node_lookup.at("mul_node1")->attr().count("bar"));
EXPECT_EQ(0, node_lookup.at("add_node2")->attr().count("foo"));
EXPECT_EQ(1, node_lookup.at("add_node2")->attr().count("bar"));
EXPECT_EQ(1, node_lookup.at("add_node2")->attr().count("bob"));
GraphDef targeted_result;
TransformFuncContext targeted_context;
targeted_context.input_names = {};
targeted_context.output_names = {"mul_node1"};
targeted_context.params.insert(
std::pair<std::string, std::vector<std::string>>(
{"op_name", {std::string("Mul")}}));
targeted_context.params.insert(
std::pair<std::string, std::vector<std::string>>(
{"attribute_name", {std::string("foo")}}));
TF_ASSERT_OK(
RemoveAttribute(graph_def, targeted_context, &targeted_result));
MapNamesToNodes(targeted_result, &node_lookup);
EXPECT_EQ(0, node_lookup.at("mul_node1")->attr().count("foo"));
EXPECT_EQ(1, node_lookup.at("mul_node1")->attr().count("bar"));
EXPECT_EQ(1, node_lookup.at("add_node2")->attr().count("foo"));
EXPECT_EQ(1, node_lookup.at("add_node2")->attr().count("bar"));
EXPECT_EQ(1, node_lookup.at("add_node2")->attr().count("bob"));
}
};
TEST_F(RemoveAttributeTest, TestRemoveAttribute) { TestRemoveAttribute(); }
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/remove_control_dependencies.cc
+46
View File
@@ -0,0 +1,46 @@
/* 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.
==============================================================================*/
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/graph/node_builder.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
// Remove control dependencies in preparation for inference.
// In the tensorflow graph, control dependencies are represented as extra
// inputs which are referenced with "^tensor_name".
// See node_def.proto for more details.
absl::Status RemoveControlDependencies(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def) {
output_graph_def->Clear();
for (const NodeDef& node : input_graph_def.node()) {
NodeDef* new_node = output_graph_def->mutable_node()->Add();
*new_node = node;
new_node->clear_input();
for (const auto& input : node.input()) {
if (input[0] != '^') {
new_node->add_input(input);
}
}
}
return absl::OkStatus();
}
REGISTER_GRAPH_TRANSFORM("remove_control_dependencies", RemoveControlDependencies);
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/remove_device.cc
+45
View File
@@ -0,0 +1,45 @@
/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/core/common_runtime/constant_folding.h"
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/graph/node_builder.h"
#include "tensorflow/core/graph/subgraph.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/tools/graph_transforms/fold_constants_lib.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
// Clears the device field of all ops in the graph.
absl::Status RemoveDevice(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def) {
output_graph_def->Clear();
for (const NodeDef& node : input_graph_def.node()) {
NodeDef* new_node = output_graph_def->mutable_node()->Add();
*new_node = node;
new_node->set_device("");
}
return absl::OkStatus();
}
REGISTER_GRAPH_TRANSFORM("remove_device", RemoveDevice);
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/remove_device_test.cc
+95
View File
@@ -0,0 +1,95 @@
/* Copyright 2015 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/cc/ops/const_op.h"
#include "tensorflow/cc/ops/image_ops.h"
#include "tensorflow/cc/ops/nn_ops.h"
#include "tensorflow/cc/ops/sendrecv_ops.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
// Declare here, so we don't need a public header.
absl::Status RemoveDevice(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def);
class RemoveDeviceTest : public ::testing::Test {
protected:
void TestRemoveDevice() {
GraphDef graph_def;
NodeDef* mul_node1 = graph_def.add_node();
mul_node1->set_name("mul_node1");
mul_node1->set_op("Mul");
mul_node1->set_device("//cpu:0");
mul_node1->add_input("add_node2");
mul_node1->add_input("add_node3");
NodeDef* add_node2 = graph_def.add_node();
add_node2->set_name("add_node2");
add_node2->set_op("Add");
add_node2->add_input("const_node1");
add_node2->add_input("const_node2");
add_node2->set_device("//device:GPU:1");
NodeDef* add_node3 = graph_def.add_node();
add_node3->set_name("add_node3");
add_node3->set_op("Add");
add_node3->add_input("const_node1");
add_node3->add_input("const_node3");
NodeDef* const_node1 = graph_def.add_node();
const_node1->set_name("const_node1");
const_node1->set_op("Const");
NodeDef* const_node2 = graph_def.add_node();
const_node2->set_name("const_node2");
const_node2->set_op("Const");
NodeDef* const_node3 = graph_def.add_node();
const_node3->set_name("const_node3");
const_node3->set_op("Const");
NodeDef* add_node4 = graph_def.add_node();
add_node4->set_name("add_node4");
add_node4->set_op("Add");
add_node4->add_input("add_node2");
add_node4->add_input("add_node3");
GraphDef result;
TransformFuncContext context;
context.input_names = {};
context.output_names = {"mul_node1"};
TF_ASSERT_OK(RemoveDevice(graph_def, context, &result));
std::map<std::string, const NodeDef*> node_lookup;
MapNamesToNodes(result, &node_lookup);
EXPECT_EQ("", node_lookup.at("mul_node1")->device());
EXPECT_EQ("", node_lookup.at("add_node2")->device());
}
};
TEST_F(RemoveDeviceTest, TestRemoveDevice) { TestRemoveDevice(); }
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/remove_nodes.cc
+116
View File
@@ -0,0 +1,116 @@
/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/core/common_runtime/constant_folding.h"
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/graph/node_builder.h"
#include "tensorflow/core/graph/subgraph.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/tools/graph_transforms/fold_constants_lib.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
// Deletes any specified types of nodes, unless they're necessary for the
// graph's inputs or outputs.
absl::Status RemoveNodes(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def) {
if (!context.params.count("op")) {
return absl::InvalidArgumentError(
"remove_nodes expects at least one 'op'"
"argument, e.g. remove_nodes(op=Identity)");
}
int32_t max_inputs;
TF_RETURN_IF_ERROR(
context.GetOneInt32Parameter("max_inputs", 1, &max_inputs));
// Make sure we don't get rid of any nodes used as graph inputs or outputs.
std::set<std::string> required_nodes;
for (const std::string& input : context.input_names) {
required_nodes.insert(NodeNameFromInput(input));
}
for (const std::string& output : context.output_names) {
required_nodes.insert(NodeNameFromInput(output));
}
std::vector<std::string> ops_to_remove = context.params.at("op");
GraphDef current_graph_def = input_graph_def;
for (const std::string& op : ops_to_remove) {
for (int num_inputs = 1; num_inputs <= max_inputs; ++num_inputs) {
// Look for a variable number of inputs.
OpTypePattern pattern = {op};
pattern.inputs.resize(num_inputs);
for (int i = 0; i < num_inputs; ++i) {
pattern.inputs[i] = {"*"};
}
// Keep looking for nodes to remove until there are no more changes.
bool any_nodes_removed;
do {
any_nodes_removed = false;
std::map<std::string, std::string> inputs_to_rename;
GraphDef replaced_graph_def;
TF_RETURN_IF_ERROR(ReplaceMatchingOpTypes(
current_graph_def, pattern,
[&inputs_to_rename, &required_nodes, &any_nodes_removed](
const NodeMatch& match,
const std::set<std::string>& input_nodes,
const std::set<std::string>& output_nodes,
std::vector<NodeDef>* new_nodes) {
const NodeDef& replace_node = match.node;
// If this node is needed in the inputs or outputs don't replace
// it.
if (required_nodes.count(replace_node.name())) {
LOG(INFO) << "Skipping replacement for " << replace_node.name();
CopyOriginalMatch(match, new_nodes);
return absl::OkStatus();
}
const NodeDef& input_node = match.inputs[0].node;
std::string target_name = input_node.name();
for (const std::string& input : replace_node.input()) {
if (!input.compare(0, target_name.size(), target_name)) {
if (input.size() == target_name.size() ||
input[target_name.size()] == ':') {
target_name = input;
break;
}
}
}
inputs_to_rename[replace_node.name()] = target_name;
inputs_to_rename["^" + replace_node.name()] =
"^" + input_node.name();
new_nodes->push_back(input_node);
any_nodes_removed = true;
return absl::OkStatus();
},
{true}, &replaced_graph_def));
// Make sure all references to removed nodes now point to their inputs.
TF_RETURN_IF_ERROR(RenameNodeInputs(
replaced_graph_def, inputs_to_rename,
std::unordered_set<std::string>(), &current_graph_def));
} while (any_nodes_removed);
}
}
*output_graph_def = current_graph_def;
return absl::OkStatus();
}
REGISTER_GRAPH_TRANSFORM("remove_nodes", RemoveNodes);
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/remove_nodes_test.cc
+278
View File
@@ -0,0 +1,278 @@
/* Copyright 2015 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/cc/ops/const_op.h"
#include "tensorflow/cc/ops/image_ops.h"
#include "tensorflow/cc/ops/nn_ops.h"
#include "tensorflow/cc/ops/sendrecv_ops.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
// Declare here, so we don't need a public header.
absl::Status RemoveNodes(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def);
class RemoveNodesTest : public ::testing::Test {
protected:
void TestRemoveNodes() {
GraphDef graph_def;
NodeDef* add_node1 = graph_def.add_node();
add_node1->set_name("add_node1");
add_node1->set_op("Add");
add_node1->add_input("add_node2");
add_node1->add_input("add_node3");
NodeDef* add_node2 = graph_def.add_node();
add_node2->set_name("add_node2");
add_node2->set_op("Add");
add_node2->add_input("identity_node1");
add_node2->add_input("identity_node2");
NodeDef* add_node3 = graph_def.add_node();
add_node3->set_name("add_node3");
add_node3->set_op("Add");
add_node3->add_input("identity_node1");
add_node3->add_input("const_node3");
NodeDef* identity_node1 = graph_def.add_node();
identity_node1->set_name("identity_node1");
identity_node1->set_op("Identity");
identity_node1->add_input("const_node1");
NodeDef* identity_node2 = graph_def.add_node();
identity_node2->set_name("identity_node2");
identity_node2->set_op("Identity");
identity_node2->add_input("const_node2");
NodeDef* identity_node3 = graph_def.add_node();
identity_node3->set_name("identity_node3");
identity_node3->set_op("Identity");
identity_node3->add_input("const_node3");
NodeDef* const_node1 = graph_def.add_node();
const_node1->set_name("const_node1");
const_node1->set_op("Const");
NodeDef* const_node2 = graph_def.add_node();
const_node2->set_name("const_node2");
const_node2->set_op("Const");
NodeDef* const_node3 = graph_def.add_node();
const_node3->set_name("const_node3");
const_node3->set_op("Const");
NodeDef* add_node4 = graph_def.add_node();
add_node4->set_name("add_node4");
add_node4->set_op("Add");
add_node4->add_input("add_node2");
add_node4->add_input("add_node3");
GraphDef result;
TransformFuncContext context;
context.input_names = {};
context.output_names = {"add_node1"};
context.params.insert(std::pair<std::string, std::vector<std::string>>(
{"op", {std::string("Identity")}}));
TF_ASSERT_OK(RemoveNodes(graph_def, context, &result));
std::map<std::string, const NodeDef*> node_lookup;
MapNamesToNodes(result, &node_lookup);
EXPECT_EQ(1, node_lookup.count("add_node1"));
EXPECT_EQ("add_node2", node_lookup.at("add_node1")->input(0));
EXPECT_EQ("add_node3", node_lookup.at("add_node1")->input(1));
EXPECT_EQ(1, node_lookup.count("add_node2"));
EXPECT_EQ("const_node1", node_lookup.at("add_node2")->input(0));
EXPECT_EQ("const_node2", node_lookup.at("add_node2")->input(1));
EXPECT_EQ(1, node_lookup.count("add_node3"));
EXPECT_EQ("const_node1", node_lookup.at("add_node3")->input(0));
EXPECT_EQ("const_node3", node_lookup.at("add_node3")->input(1));
EXPECT_EQ(1, node_lookup.count("add_node4"));
EXPECT_EQ("add_node2", node_lookup.at("add_node4")->input(0));
EXPECT_EQ("add_node3", node_lookup.at("add_node4")->input(1));
EXPECT_EQ(0, node_lookup.count("identity_node1"));
EXPECT_EQ(0, node_lookup.count("identity_node2"));
EXPECT_EQ(0, node_lookup.count("identity_node3"));
EXPECT_EQ(1, node_lookup.count("const_node1"));
EXPECT_EQ("Const", node_lookup.at("const_node1")->op());
EXPECT_EQ(1, node_lookup.count("const_node2"));
EXPECT_EQ("Const", node_lookup.at("const_node2")->op());
EXPECT_EQ(1, node_lookup.count("const_node3"));
EXPECT_EQ("Const", node_lookup.at("const_node3")->op());
}
void TestRemoveOutputNodes() {
GraphDef graph_def;
NodeDef* const_node1 = graph_def.add_node();
const_node1->set_name("const_node1");
const_node1->set_op("Const");
NodeDef* const_node2 = graph_def.add_node();
const_node2->set_name("const_node2");
const_node2->set_op("Const");
NodeDef* add_node = graph_def.add_node();
add_node->set_name("add_node");
add_node->set_op("Add");
add_node->add_input("const_node1");
add_node->add_input("const_node2");
NodeDef* identity_node = graph_def.add_node();
identity_node->set_name("identity_node");
identity_node->set_op("Identity");
identity_node->add_input("add_node");
GraphDef result;
TransformFuncContext context;
context.input_names = {};
context.output_names = {"identity_node"};
context.params.insert(std::pair<std::string, std::vector<std::string>>(
{"op", {std::string("Identity")}}));
TF_ASSERT_OK(RemoveNodes(graph_def, context, &result));
std::map<std::string, const NodeDef*> node_lookup;
MapNamesToNodes(result, &node_lookup);
EXPECT_EQ(1, node_lookup.count("add_node"));
EXPECT_EQ("const_node1", node_lookup.at("add_node")->input(0));
EXPECT_EQ("const_node2", node_lookup.at("add_node")->input(1));
EXPECT_EQ(1, node_lookup.count("identity_node"));
EXPECT_EQ("add_node", node_lookup.at("identity_node")->input(0));
}
void TestRemoveChainedNodes() {
GraphDef graph_def;
NodeDef* const_node1 = graph_def.add_node();
const_node1->set_name("const_node1");
const_node1->set_op("Const");
NodeDef* identity_node1 = graph_def.add_node();
identity_node1->set_name("identity_node1");
identity_node1->set_op("Identity");
identity_node1->add_input("const_node1");
NodeDef* identity_node2 = graph_def.add_node();
identity_node2->set_name("identity_node2");
identity_node2->set_op("Identity");
identity_node2->add_input("identity_node1");
NodeDef* identity_node3 = graph_def.add_node();
identity_node3->set_name("identity_node3");
identity_node3->set_op("Identity");
identity_node3->add_input("identity_node2");
NodeDef* const_node2 = graph_def.add_node();
const_node2->set_name("const_node2");
const_node2->set_op("Const");
NodeDef* add_node = graph_def.add_node();
add_node->set_name("add_node");
add_node->set_op("Add");
add_node->add_input("identity_node3");
add_node->add_input("const_node2");
GraphDef result;
TransformFuncContext context;
context.input_names = {};
context.output_names = {"identity_node"};
context.params.insert(std::pair<std::string, std::vector<std::string>>(
{"op", {std::string("Identity")}}));
TF_ASSERT_OK(RemoveNodes(graph_def, context, &result));
std::map<std::string, const NodeDef*> node_lookup;
MapNamesToNodes(result, &node_lookup);
EXPECT_EQ(1, node_lookup.count("add_node"));
EXPECT_EQ("const_node1", node_lookup.at("add_node")->input(0));
EXPECT_EQ("const_node2", node_lookup.at("add_node")->input(1));
EXPECT_EQ(0, node_lookup.count("identity_node1"));
EXPECT_EQ(0, node_lookup.count("identity_node2"));
EXPECT_EQ(0, node_lookup.count("identity_node3"));
}
void TestRemoveMultipleInputs() {
GraphDef graph_def;
NodeDef* const_node1 = graph_def.add_node();
const_node1->set_name("const_node1");
const_node1->set_op("Const");
NodeDef* const_node2 = graph_def.add_node();
const_node2->set_name("const_node2");
const_node2->set_op("Const");
NodeDef* const_node3 = graph_def.add_node();
const_node3->set_name("const_node3");
const_node3->set_op("Const");
NodeDef* const_node4 = graph_def.add_node();
const_node4->set_name("const_node4");
const_node4->set_op("Const");
NodeDef* fake_quant_node = graph_def.add_node();
fake_quant_node->set_name("fake_quant_node");
fake_quant_node->set_op("FakeQuantWithMinMaxVars");
fake_quant_node->add_input("const_node1");
fake_quant_node->add_input("const_node2");
fake_quant_node->add_input("const_node3");
NodeDef* add_node = graph_def.add_node();
add_node->set_name("add_node");
add_node->set_op("Add");
add_node->add_input("fake_quant_node");
add_node->add_input("const_node4");
GraphDef result;
TransformFuncContext context;
context.input_names = {};
context.output_names = {"add_node"};
context.params.insert(std::pair<std::string, std::vector<std::string>>(
{"op", {std::string("FakeQuantWithMinMaxVars")}}));
context.params.insert(std::pair<std::string, std::vector<std::string>>(
{"max_inputs", {std::string("3")}}));
TF_ASSERT_OK(RemoveNodes(graph_def, context, &result));
std::map<std::string, const NodeDef*> node_lookup;
MapNamesToNodes(result, &node_lookup);
ASSERT_EQ(1, node_lookup.count("const_node1"));
ASSERT_EQ(1, node_lookup.count("const_node4"));
ASSERT_EQ(0, node_lookup.count("fake_quant_node"));
ASSERT_EQ(1, node_lookup.count("add_node"));
EXPECT_EQ("const_node1", node_lookup.at("add_node")->input(0));
EXPECT_EQ("const_node4", node_lookup.at("add_node")->input(1));
}
};
TEST_F(RemoveNodesTest, TestRemoveNodes) { TestRemoveNodes(); }
TEST_F(RemoveNodesTest, TestRemoveOutputNodes) { TestRemoveOutputNodes(); }
TEST_F(RemoveNodesTest, TestRemoveChainedNodes) { TestRemoveChainedNodes(); }
TEST_F(RemoveNodesTest, TestRemoveMultipleInputs) {
TestRemoveMultipleInputs();
}
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/rename_attribute.cc
+70
View File
@@ -0,0 +1,70 @@
/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/core/common_runtime/constant_folding.h"
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/graph/node_builder.h"
#include "tensorflow/core/graph/subgraph.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/tools/graph_transforms/fold_constants_lib.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
absl::Status RenameAttribute(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def) {
if (!context.params.count("old_attribute_name") ||
(context.params.at("old_attribute_name").size() != 1) ||
!context.params.count("new_attribute_name") ||
(context.params.at("new_attribute_name").size() != 1)) {
return absl::InvalidArgumentError(
"rename_attribute expects exactly one 'old_attribute_name' and one "
"'new_attribute_name' argument, e.g. "
"rename_attribute(old_attribute_name=foo, new_attribute_name=bar)");
}
std::string op_name;
if (context.params.count("op_name")) {
op_name = context.params.at("op_name")[0];
} else {
op_name = "*";
}
const std::string old_attribute_name =
context.params.at("old_attribute_name")[0];
const std::string new_attribute_name =
context.params.at("new_attribute_name")[0];
output_graph_def->Clear();
for (const NodeDef& node : input_graph_def.node()) {
NodeDef* new_node = output_graph_def->mutable_node()->Add();
*new_node = node;
if (((op_name == "*") || (op_name == node.op())) &&
(node.attr().count(old_attribute_name))) {
AttrValue attribute_value = node.attr().at(old_attribute_name);
new_node->mutable_attr()->erase(old_attribute_name);
new_node->mutable_attr()->insert({new_attribute_name, attribute_value});
}
}
return absl::OkStatus();
}
REGISTER_GRAPH_TRANSFORM("rename_attribute", RenameAttribute);
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/rename_attribute_test.cc
+134
View File
@@ -0,0 +1,134 @@
/* Copyright 2015 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/cc/ops/const_op.h"
#include "tensorflow/cc/ops/image_ops.h"
#include "tensorflow/cc/ops/nn_ops.h"
#include "tensorflow/cc/ops/sendrecv_ops.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
// Declare here, so we don't need a public header.
absl::Status RenameAttribute(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def);
class RenameAttributeTest : public ::testing::Test {
protected:
void TestRenameAttribute() {
GraphDef graph_def;
NodeDef* mul_node1 = graph_def.add_node();
mul_node1->set_name("mul_node1");
mul_node1->set_op("Mul");
mul_node1->add_input("add_node2");
mul_node1->add_input("add_node3");
AddNodeAttr("foo", 23, mul_node1);
AddNodeAttr("bar", "something", mul_node1);
NodeDef* add_node2 = graph_def.add_node();
add_node2->set_name("add_node2");
add_node2->set_op("Add");
add_node2->add_input("const_node1");
add_node2->add_input("const_node2");
AddNodeAttr("foo", 46, add_node2);
AddNodeAttr("bob", 23, add_node2);
AddNodeAttr("bar", "something else", add_node2);
NodeDef* add_node3 = graph_def.add_node();
add_node3->set_name("add_node3");
add_node3->set_op("Add");
add_node3->add_input("const_node1");
add_node3->add_input("const_node3");
NodeDef* const_node1 = graph_def.add_node();
const_node1->set_name("const_node1");
const_node1->set_op("Const");
NodeDef* const_node2 = graph_def.add_node();
const_node2->set_name("const_node2");
const_node2->set_op("Const");
NodeDef* const_node3 = graph_def.add_node();
const_node3->set_name("const_node3");
const_node3->set_op("Const");
NodeDef* add_node4 = graph_def.add_node();
add_node4->set_name("add_node4");
add_node4->set_op("Add");
add_node4->add_input("add_node2");
add_node4->add_input("add_node3");
GraphDef wildcard_result;
TransformFuncContext context;
context.input_names = {};
context.output_names = {"mul_node1"};
context.params.insert(std::pair<std::string, std::vector<std::string>>(
{"op_name", {std::string("*")}}));
context.params.insert(std::pair<std::string, std::vector<std::string>>(
{"old_attribute_name", {std::string("foo")}}));
context.params.insert(std::pair<std::string, std::vector<std::string>>(
{"new_attribute_name", {std::string("baz")}}));
TF_ASSERT_OK(RenameAttribute(graph_def, context, &wildcard_result));
std::map<std::string, const NodeDef*> node_lookup;
MapNamesToNodes(wildcard_result, &node_lookup);
EXPECT_EQ(0, node_lookup.at("mul_node1")->attr().count("foo"));
EXPECT_EQ(1, node_lookup.at("mul_node1")->attr().count("baz"));
EXPECT_EQ(1, node_lookup.at("mul_node1")->attr().count("bar"));
EXPECT_EQ(0, node_lookup.at("add_node2")->attr().count("foo"));
EXPECT_EQ(1, node_lookup.at("add_node2")->attr().count("baz"));
EXPECT_EQ(1, node_lookup.at("add_node2")->attr().count("bar"));
EXPECT_EQ(1, node_lookup.at("add_node2")->attr().count("bob"));
GraphDef targeted_result;
TransformFuncContext targeted_context;
targeted_context.input_names = {};
targeted_context.output_names = {"mul_node1"};
targeted_context.params.insert(
std::pair<std::string, std::vector<std::string>>(
{"op_name", {std::string("Mul")}}));
targeted_context.params.insert(
std::pair<std::string, std::vector<std::string>>(
{"old_attribute_name", {std::string("foo")}}));
targeted_context.params.insert(
std::pair<std::string, std::vector<std::string>>(
{"new_attribute_name", {std::string("baz")}}));
TF_ASSERT_OK(
RenameAttribute(graph_def, targeted_context, &targeted_result));
MapNamesToNodes(targeted_result, &node_lookup);
EXPECT_EQ(0, node_lookup.at("mul_node1")->attr().count("foo"));
EXPECT_EQ(1, node_lookup.at("mul_node1")->attr().count("baz"));
EXPECT_EQ(1, node_lookup.at("mul_node1")->attr().count("bar"));
EXPECT_EQ(1, node_lookup.at("add_node2")->attr().count("foo"));
EXPECT_EQ(0, node_lookup.at("add_node2")->attr().count("baz"));
EXPECT_EQ(1, node_lookup.at("add_node2")->attr().count("bar"));
EXPECT_EQ(1, node_lookup.at("add_node2")->attr().count("bob"));
}
};
TEST_F(RenameAttributeTest, TestRenameAttribute) { TestRenameAttribute(); }
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/rename_node.cc
+71
View File
@@ -0,0 +1,71 @@
/* 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 <string>
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
absl::Status RenameNode(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def) {
if (!context.params.count("old_node_name") ||
(context.params.at("old_node_name").size() != 1) ||
!context.params.count("new_node_name") ||
(context.params.at("new_node_name").size() != 1)) {
return absl::InvalidArgumentError(
"rename_node expects exactly one 'old_node_name' and one "
"'new_node_name' argument, e.g. "
"rename_node(old_attribute_name=super/deep/output, "
"new_attribute_name=output)");
}
const std::string old_node_name = context.params.at("old_node_name")[0];
const std::string new_node_name = context.params.at("new_node_name")[0];
output_graph_def->Clear();
for (const NodeDef& input_node : input_graph_def.node()) {
NodeDef* node = output_graph_def->mutable_node()->Add();
*node = input_node;
if (node->name() == new_node_name) {
return absl::Status(
absl::StatusCode::kInvalidArgument,
"A node is alreading using " + new_node_name + "as name.");
}
if (node->name() == old_node_name) {
node->set_name(new_node_name);
}
for (std::string& input_name : *node->mutable_input()) {
std::string prefix;
std::string input_node_name;
std::string suffix;
NodeNamePartsFromInput(input_name, &prefix, &input_node_name, &suffix);
if (input_node_name == old_node_name) {
std::string new_input_name = prefix + new_node_name + suffix;
input_name = new_input_name;
}
}
}
return absl::OkStatus();
}
REGISTER_GRAPH_TRANSFORM("rename_node", RenameNode);
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/rename_node_test.cc
+96
View File
@@ -0,0 +1,96 @@
/* 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 <string>
#include <utility>
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
// Declare here, so we don't need a public header.
absl::Status RenameNode(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def);
TEST(RenameNodeTest, Rename) {
GraphDef in_graph;
NodeDef* node = in_graph.add_node();
node->set_name("input");
node->set_op("Placeholder");
NodeDef* node_splitter = in_graph.add_node();
node_splitter->set_name("splitter");
node_splitter->set_op("Split");
NodeDef* node_adder = in_graph.add_node();
node_adder->set_op("Add");
node_adder->set_name("adder");
node_adder->add_input("splitter");
node_adder->add_input("splitter:1");
GraphDef result;
TransformFuncContext context;
context.input_names = {};
context.output_names = {"adder"};
context.params.insert(std::pair<std::string, std::vector<std::string>>(
{"old_node_name", {std::string("splitter")}}));
context.params.insert(std::pair<std::string, std::vector<std::string>>(
{"new_node_name", {std::string("demux")}}));
TF_ASSERT_OK(RenameNode(in_graph, context, &result));
std::map<std::string, const NodeDef*> node_lookup;
MapNamesToNodes(result, &node_lookup);
EXPECT_EQ(1, node_lookup.count("demux"));
EXPECT_EQ(1, node_lookup.count("adder"));
EXPECT_EQ(2, node_lookup["adder"]->input().size());
EXPECT_EQ("demux", node_lookup["adder"]->input()[0]);
EXPECT_EQ("demux:1", node_lookup["adder"]->input()[1]);
}
TEST(RenameNodeTest, FailWhenNameAlreadyExists) {
GraphDef in_graph;
NodeDef* node = in_graph.add_node();
node->set_name("input");
node->set_op("Placeholder");
NodeDef* node_splitter = in_graph.add_node();
node_splitter->set_name("splitter");
node_splitter->set_op("Split");
NodeDef* node_adder = in_graph.add_node();
node_adder->set_op("Add");
node_adder->set_name("adder");
node_adder->add_input("splitter");
node_adder->add_input("splitter:1");
GraphDef result;
TransformFuncContext context;
context.input_names = {};
context.output_names = {"adder"};
context.params.insert(std::pair<std::string, std::vector<std::string>>(
{"old_node_name", {std::string("splitter")}}));
context.params.insert(std::pair<std::string, std::vector<std::string>>(
{"new_node_name", {std::string("adder")}}));
EXPECT_FALSE(RenameNode(in_graph, context, &result).ok());
}
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/rename_op.cc
+58
View File
@@ -0,0 +1,58 @@
/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/core/common_runtime/constant_folding.h"
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/graph/node_builder.h"
#include "tensorflow/core/graph/subgraph.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/tools/graph_transforms/fold_constants_lib.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
// Changes the op type of a specified op.
absl::Status RenameOp(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def) {
if (!context.params.count("old_op_name") ||
(context.params.at("old_op_name").size() != 1) ||
!context.params.count("new_op_name") ||
(context.params.at("new_op_name").size() != 1)) {
return absl::InvalidArgumentError(
"rename_op expects exactly one 'old_op_name' and 'new_op_name' "
"argument, e.g. rename_op(old_op_name=Mul, new_op_name=Multiply)");
}
const std::string old_op_name = context.params.at("old_op_name")[0];
const std::string new_op_name = context.params.at("new_op_name")[0];
output_graph_def->Clear();
for (const NodeDef& node : input_graph_def.node()) {
NodeDef* new_node = output_graph_def->mutable_node()->Add();
*new_node = node;
if (node.op() == old_op_name) {
new_node->set_op(new_op_name);
}
}
return absl::OkStatus();
}
REGISTER_GRAPH_TRANSFORM("rename_op", RenameOp);
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/rename_op_test.cc
+109
View File
@@ -0,0 +1,109 @@
/* Copyright 2015 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/cc/ops/const_op.h"
#include "tensorflow/cc/ops/image_ops.h"
#include "tensorflow/cc/ops/nn_ops.h"
#include "tensorflow/cc/ops/sendrecv_ops.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
// Declare here, so we don't need a public header.
absl::Status RenameOp(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def);
class RenameOpTest : public ::testing::Test {
protected:
void TestRenameOp() {
GraphDef graph_def;
NodeDef* mul_node1 = graph_def.add_node();
mul_node1->set_name("mul_node1");
mul_node1->set_op("Mul");
mul_node1->add_input("add_node2");
mul_node1->add_input("add_node3");
NodeDef* add_node2 = graph_def.add_node();
add_node2->set_name("add_node2");
add_node2->set_op("Add");
add_node2->add_input("const_node1");
add_node2->add_input("const_node2");
NodeDef* add_node3 = graph_def.add_node();
add_node3->set_name("add_node3");
add_node3->set_op("Add");
add_node3->add_input("const_node1");
add_node3->add_input("const_node3");
NodeDef* const_node1 = graph_def.add_node();
const_node1->set_name("const_node1");
const_node1->set_op("Const");
NodeDef* const_node2 = graph_def.add_node();
const_node2->set_name("const_node2");
const_node2->set_op("Const");
NodeDef* const_node3 = graph_def.add_node();
const_node3->set_name("const_node3");
const_node3->set_op("Const");
NodeDef* add_node4 = graph_def.add_node();
add_node4->set_name("add_node4");
add_node4->set_op("Add");
add_node4->add_input("add_node2");
add_node4->add_input("add_node3");
GraphDef result;
TransformFuncContext context;
context.input_names = {};
context.output_names = {"mul_node1"};
context.params.insert(std::pair<std::string, std::vector<std::string>>(
{"old_op_name", {std::string("Mul")}}));
context.params.insert(std::pair<std::string, std::vector<std::string>>(
{"new_op_name", {std::string("Multiply")}}));
TF_ASSERT_OK(RenameOp(graph_def, context, &result));
std::map<std::string, const NodeDef*> node_lookup;
MapNamesToNodes(result, &node_lookup);
EXPECT_EQ(1, node_lookup.count("mul_node1"));
EXPECT_EQ("Multiply", node_lookup.at("mul_node1")->op());
EXPECT_EQ(1, node_lookup.count("add_node2"));
EXPECT_EQ("Add", node_lookup.at("add_node2")->op());
EXPECT_EQ(1, node_lookup.count("add_node3"));
EXPECT_EQ("Add", node_lookup.at("add_node3")->op());
EXPECT_EQ(1, node_lookup.count("add_node4"));
EXPECT_EQ("Add", node_lookup.at("add_node4")->op());
EXPECT_EQ(1, node_lookup.count("const_node1"));
EXPECT_EQ("Const", node_lookup.at("const_node1")->op());
EXPECT_EQ(1, node_lookup.count("const_node2"));
EXPECT_EQ("Const", node_lookup.at("const_node2")->op());
EXPECT_EQ(1, node_lookup.count("const_node3"));
EXPECT_EQ("Const", node_lookup.at("const_node3")->op());
}
};
TEST_F(RenameOpTest, TestRenameOp) { TestRenameOp(); }
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/round_weights.cc
+117
View File
@@ -0,0 +1,117 @@
/* Copyright 2016 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.
==============================================================================*/
#define EIGEN_USE_THREADS
#include "tensorflow/core/common_runtime/constant_folding.h"
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/common_runtime/threadpool_device.h"
#include "tensorflow/core/graph/node_builder.h"
#include "tensorflow/core/graph/subgraph.h"
#include "tensorflow/core/kernels/quantization_utils.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
// Rounds any large float constants to the specified number of levels.
absl::Status RoundWeights(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def) {
int32_t num_steps;
TF_RETURN_IF_ERROR(
context.GetOneInt32Parameter("num_steps", 256, &num_steps));
TF_RETURN_IF_ERROR(ReplaceMatchingOpTypes(
input_graph_def, {"Const"},
[num_steps](const NodeMatch& match,
const std::set<std::string>& input_nodes,
const std::set<std::string>& output_nodes,
std::vector<NodeDef>* new_nodes) {
const NodeDef& old_const_node = match.node;
if (!old_const_node.attr().count("dtype")) {
return absl::InvalidArgumentError(absl::StrCat(
"No 'dtype' attribute for Const node ", old_const_node.name()));
}
if (!old_const_node.attr().count("value")) {
return absl::InvalidArgumentError(absl::StrCat(
"No 'value' attribute for Const node ", old_const_node.name()));
}
const DataType old_dtype = old_const_node.attr().at("dtype").type();
Tensor old_tensor;
if (!old_tensor.FromProto(old_const_node.attr().at("value").tensor())) {
return absl::InvalidArgumentError(absl::StrCat(
"Decoding Tensor failed for node", old_const_node.name()));
}
const size_t num_elements = old_tensor.NumElements();
// If this isn't a float constant, or it's too small, then reuse the
// same node with no changes. The size is important because small
// constants tend to be used for more accuracy-sensitive calculations,
// and the benefit of shrinking them is very marginal.
if ((old_dtype != DT_FLOAT) || (num_elements < 16)) {
new_nodes->push_back(old_const_node);
return absl::OkStatus();
}
const float* old_values = old_tensor.flat<float>().data();
float min = std::numeric_limits<float>::max();
float max = std::numeric_limits<float>::min();
for (int i = 0; i < num_elements; ++i) {
const float value = old_values[i];
min = std::min(min, value);
max = std::max(max, value);
}
// min_value == max_value is a tricky case. It can occur for general
// tensors, and of course for scalars. The quantized ops cannot deal
// with this case, so we set max_value to something else.
// It's a tricky question what is the numerically best solution to
// deal with this degeneracy.
// TODO(petewarden): Better use a tolerance than a hard comparison?
if (min == max) {
if (std::abs(min) < 0.000001f) {
max = min + 1.0f;
} else if (min > 0) {
max = 2.0f * min;
} else {
min = 2.0f * max;
}
}
Tensor rounded_tensor(DT_FLOAT, old_tensor.shape());
float* rounded_values = rounded_tensor.flat<float>().data();
const float bucket_width = (max - min) / num_steps;
for (int i = 0; i < num_elements; ++i) {
const int32_t bucket =
std::floor((old_values[i] - min) / bucket_width);
rounded_values[i] = min + (bucket_width * (bucket + 0.5f));
}
NodeDef rounded_const_node;
rounded_const_node.set_op("Const");
rounded_const_node.set_name(old_const_node.name());
SetNodeAttr("dtype", DT_FLOAT, &rounded_const_node);
SetNodeTensorAttr<float>("value", rounded_tensor, &rounded_const_node);
new_nodes->push_back(rounded_const_node);
return absl::OkStatus();
},
{}, output_graph_def));
return absl::OkStatus();
}
REGISTER_GRAPH_TRANSFORM("round_weights", RoundWeights);
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/round_weights_test.cc
+96
View File
@@ -0,0 +1,96 @@
/* Copyright 2015 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/cc/ops/const_op.h"
#include "tensorflow/cc/ops/image_ops.h"
#include "tensorflow/cc/ops/nn_ops.h"
#include "tensorflow/cc/ops/sendrecv_ops.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
// Declare here, so we don't need a public header.
absl::Status RoundWeights(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def);
class RoundWeightsTest : public ::testing::Test {
protected:
void TestRoundWeights() {
auto root = tensorflow::Scope::NewRootScope();
using namespace ::tensorflow::ops; // NOLINT(build/namespaces)
Tensor input_data(DT_FLOAT, TensorShape({1, 1, 6, 2}));
test::FillValues<float>(
&input_data, {1.0f, 4.0f, 2.0f, 5.0f, 3.0f, 6.0f, -1.0f, -4.0f, -2.0f,
-5.0f, -3.0f, -6.0f});
Output input_op =
Const(root.WithOpName("input_op"), Input::Initializer(input_data));
Tensor weights_data(DT_FLOAT, TensorShape({1, 2, 2, 10}));
test::FillValues<float>(
&weights_data,
{1.0f, 2.0f, 3.0f, 4.0f, 0.1f, 0.2f, 0.3f, 0.4f, 1.0f, 2.0f,
3.0f, 4.0f, 0.1f, 0.2f, 0.3f, 0.4f, 1.0f, 2.0f, 3.0f, 4.0f,
0.1f, 0.2f, 0.3f, 0.4f, 1.0f, 2.0f, 3.0f, 4.0f, 0.1f, 0.2f,
0.3f, 0.4f, 1.0f, 2.0f, 3.0f, 4.0f, 0.1f, 0.2f, 0.3f, 0.4f});
Output weights_op =
Const(root.WithOpName("weights_op"), Input::Initializer(weights_data));
Output conv_op = Conv2D(root.WithOpName("output"), input_op, weights_op,
{1, 1, 1, 1}, "VALID");
GraphDef original_graph_def;
TF_ASSERT_OK(root.ToGraphDef(&original_graph_def));
std::unique_ptr<Session> original_session(NewSession(SessionOptions()));
TF_ASSERT_OK(original_session->Create(original_graph_def));
std::vector<Tensor> original_outputs;
TF_ASSERT_OK(original_session->Run({}, {"output"}, {}, &original_outputs));
GraphDef rounded_graph_def;
TF_ASSERT_OK(
RoundWeights(original_graph_def, {{}, {"output"}}, &rounded_graph_def));
std::unique_ptr<Session> rounded_session(NewSession(SessionOptions()));
TF_ASSERT_OK(rounded_session->Create(rounded_graph_def));
std::vector<Tensor> rounded_outputs;
TF_ASSERT_OK(rounded_session->Run({}, {"output"}, {}, &rounded_outputs));
test::ExpectTensorNear<float>(original_outputs[0], rounded_outputs[0], 0.5);
std::map<std::string, const NodeDef*> node_lookup;
MapNamesToNodes(rounded_graph_def, &node_lookup);
EXPECT_EQ(1, node_lookup.count("input_op"));
const NodeDef* r_input_op = node_lookup.at("input_op");
EXPECT_EQ(DT_FLOAT, r_input_op->attr().at("dtype").type());
EXPECT_EQ(1, node_lookup.count("weights_op"));
const NodeDef* r_weights_op = node_lookup.at("weights_op");
EXPECT_EQ("Const", r_weights_op->op());
EXPECT_EQ(DT_FLOAT, r_weights_op->attr().at("dtype").type());
}
};
TEST_F(RoundWeightsTest, TestRoundWeights) { TestRoundWeights(); }
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/set_device.cc
+47
View File
@@ -0,0 +1,47 @@
/* 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/core/framework/node_def.pb.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
// Sets the device field of ops in the graph.
absl::Status SetDevice(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def) {
std::string new_device;
TF_RETURN_IF_ERROR(context.GetOneStringParameter("device", "", &new_device));
bool if_default;
TF_RETURN_IF_ERROR(
context.GetOneBoolParameter("if_default", false, &if_default));
output_graph_def->Clear();
for (const NodeDef& node : input_graph_def.node()) {
NodeDef* new_node = output_graph_def->mutable_node()->Add();
*new_node = node;
if (!if_default || (node.device().empty())) {
new_node->set_device(new_device);
}
}
return absl::OkStatus();
}
REGISTER_GRAPH_TRANSFORM("set_device", SetDevice);
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/set_device_test.cc
+127
View File
@@ -0,0 +1,127 @@
/* 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/cc/ops/const_op.h"
#include "tensorflow/cc/ops/image_ops.h"
#include "tensorflow/cc/ops/nn_ops.h"
#include "tensorflow/cc/ops/sendrecv_ops.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
// Declare here, so we don't need a public header.
absl::Status SetDevice(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def);
namespace {
GraphDef CreateDeviceGraph() {
GraphDef graph_def;
NodeDef* mul_node1 = graph_def.add_node();
mul_node1->set_name("mul_node1");
mul_node1->set_op("Mul");
mul_node1->set_device("/device:CPU:0");
mul_node1->add_input("add_node2");
mul_node1->add_input("add_node3");
NodeDef* add_node2 = graph_def.add_node();
add_node2->set_name("add_node2");
add_node2->set_op("Add");
add_node2->add_input("const_node1");
add_node2->add_input("const_node2");
add_node2->set_device("/device:GPU:1");
NodeDef* add_node3 = graph_def.add_node();
add_node3->set_name("add_node3");
add_node3->set_op("Add");
add_node3->add_input("const_node1");
add_node3->add_input("const_node3");
NodeDef* const_node1 = graph_def.add_node();
const_node1->set_name("const_node1");
const_node1->set_op("Const");
NodeDef* const_node2 = graph_def.add_node();
const_node2->set_name("const_node2");
const_node2->set_op("Const");
NodeDef* const_node3 = graph_def.add_node();
const_node3->set_name("const_node3");
const_node3->set_op("Const");
NodeDef* add_node4 = graph_def.add_node();
add_node4->set_name("add_node4");
add_node4->set_op("Add");
add_node4->add_input("add_node2");
add_node4->add_input("add_node3");
return graph_def;
}
} // namespace
TEST(SetDeviceTest, TestSetDevice) {
GraphDef graph_def = CreateDeviceGraph();
GraphDef result;
TransformFuncContext context;
context.input_names = {};
context.output_names = {"mul_node1"};
context.params.insert(std::pair<std::string, std::vector<std::string>>(
{"device", {std::string("/device:CPU:0")}}));
TF_ASSERT_OK(SetDevice(graph_def, context, &result));
std::map<std::string, const NodeDef*> node_lookup;
MapNamesToNodes(result, &node_lookup);
EXPECT_EQ("/device:CPU:0", node_lookup.at("mul_node1")->device());
EXPECT_EQ("/device:CPU:0", node_lookup.at("add_node2")->device());
EXPECT_EQ("/device:CPU:0", node_lookup.at("add_node3")->device());
EXPECT_EQ("/device:CPU:0", node_lookup.at("const_node1")->device());
EXPECT_EQ("/device:CPU:0", node_lookup.at("const_node2")->device());
EXPECT_EQ("/device:CPU:0", node_lookup.at("const_node3")->device());
EXPECT_EQ("/device:CPU:0", node_lookup.at("add_node4")->device());
}
TEST(SetDeviceTest, TestSetDeviceIfDefault) {
GraphDef graph_def = CreateDeviceGraph();
GraphDef result;
TransformFuncContext context;
context.input_names = {};
context.output_names = {"mul_node1"};
context.params.insert(std::pair<std::string, std::vector<std::string>>(
{"device", {std::string("/device:GPU:0")}}));
context.params.insert(std::pair<std::string, std::vector<std::string>>(
{"if_default", {std::string("true")}}));
TF_ASSERT_OK(SetDevice(graph_def, context, &result));
std::map<std::string, const NodeDef*> node_lookup;
MapNamesToNodes(result, &node_lookup);
EXPECT_EQ("/device:CPU:0", node_lookup.at("mul_node1")->device());
EXPECT_EQ("/device:GPU:1", node_lookup.at("add_node2")->device());
EXPECT_EQ("/device:GPU:0", node_lookup.at("add_node3")->device());
EXPECT_EQ("/device:GPU:0", node_lookup.at("const_node1")->device());
EXPECT_EQ("/device:GPU:0", node_lookup.at("const_node2")->device());
EXPECT_EQ("/device:GPU:0", node_lookup.at("const_node3")->device());
EXPECT_EQ("/device:GPU:0", node_lookup.at("add_node4")->device());
}
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/sort_by_execution_order.cc
+41
View File
@@ -0,0 +1,41 @@
/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/core/common_runtime/constant_folding.h"
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/graph/node_builder.h"
#include "tensorflow/core/graph/subgraph.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/tools/graph_transforms/fold_constants_lib.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
// This is a thin wrapper with the standard TransformFunc interface to the
// underlying utility function. The only difference is that we don't use the
// input or output name arguments.
absl::Status SortByExecutionOrderWithUnusedContext(
const GraphDef& input_graph_def, const TransformFuncContext& unused_context,
GraphDef* output_graph_def) {
return SortByExecutionOrder(input_graph_def, output_graph_def);
}
REGISTER_GRAPH_TRANSFORM("sort_by_execution_order",
SortByExecutionOrderWithUnusedContext);
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/sort_by_execution_order_test.cc
+206
View File
@@ -0,0 +1,206 @@
/* Copyright 2015 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/cc/ops/const_op.h"
#include "tensorflow/cc/ops/image_ops.h"
#include "tensorflow/cc/ops/nn_ops.h"
#include "tensorflow/cc/ops/sendrecv_ops.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
class SortByExecutionOrderTest : public ::testing::Test {
protected:
void GetOrder(const GraphDef& graph_def, std::map<std::string, int>* order) {
for (int i = 0; i < graph_def.node_size(); ++i) {
const NodeDef& node = graph_def.node(i);
(*order)[node.name()] = i;
}
}
void TestSimpleAdd() {
GraphDef graph_def;
NodeDef* add_node = graph_def.add_node();
add_node->set_name("add_node");
add_node->set_op("Add");
add_node->add_input("a_node");
add_node->add_input("b_node");
NodeDef* b_node = graph_def.add_node();
b_node->set_name("b_node");
b_node->set_op("Const");
NodeDef* a_node = graph_def.add_node();
a_node->set_name("a_node");
a_node->set_op("Const");
GraphDef result;
TF_ASSERT_OK(SortByExecutionOrder(graph_def, &result));
std::map<std::string, int> order;
GetOrder(result, &order);
EXPECT_EQ(2, order["add_node"]);
EXPECT_GT(2, order["a_node"]);
EXPECT_GT(2, order["b_node"]);
}
void TestSimpleLinear() {
GraphDef graph_def;
NodeDef* negative_node = graph_def.add_node();
negative_node->set_name("negative_node");
negative_node->set_op("Negative");
negative_node->add_input("sqrt_node");
NodeDef* relu_node = graph_def.add_node();
relu_node->set_name("relu_node");
relu_node->set_op("Relu");
relu_node->add_input("const_node");
NodeDef* sqrt_node = graph_def.add_node();
sqrt_node->set_name("sqrt_node");
sqrt_node->set_op("Sqrt");
sqrt_node->add_input("relu_node");
NodeDef* const_node = graph_def.add_node();
const_node->set_name("const_node");
const_node->set_op("Const");
GraphDef result;
TF_ASSERT_OK(SortByExecutionOrder(graph_def, &result));
std::map<std::string, int> order;
GetOrder(result, &order);
EXPECT_EQ(3, order["negative_node"]);
EXPECT_EQ(2, order["sqrt_node"]);
EXPECT_EQ(1, order["relu_node"]);
EXPECT_EQ(0, order["const_node"]);
}
void TestSimpleTree() {
GraphDef graph_def;
NodeDef* add_node1 = graph_def.add_node();
add_node1->set_name("add_node1");
add_node1->set_op("Add");
add_node1->add_input("add_node2");
add_node1->add_input("add_node3");
NodeDef* add_node2 = graph_def.add_node();
add_node2->set_name("add_node2");
add_node2->set_op("Add");
add_node2->add_input("const_node1");
add_node2->add_input("const_node2");
NodeDef* add_node3 = graph_def.add_node();
add_node3->set_name("add_node3");
add_node3->set_op("Add");
add_node3->add_input("const_node3");
add_node3->add_input("const_node4");
NodeDef* const_node1 = graph_def.add_node();
const_node1->set_name("const_node1");
const_node1->set_op("Const");
NodeDef* const_node2 = graph_def.add_node();
const_node2->set_name("const_node2");
const_node2->set_op("Const");
NodeDef* const_node3 = graph_def.add_node();
const_node3->set_name("const_node3");
const_node3->set_op("Const");
NodeDef* const_node4 = graph_def.add_node();
const_node4->set_name("const_node4");
const_node4->set_op("Const");
GraphDef result;
TF_ASSERT_OK(SortByExecutionOrder(graph_def, &result));
std::map<std::string, int> order;
GetOrder(result, &order);
EXPECT_EQ(6, order["add_node1"]);
EXPECT_GT(6, order["add_node2"]);
EXPECT_GT(6, order["add_node3"]);
EXPECT_GT(5, order["const_node1"]);
EXPECT_GT(5, order["const_node2"]);
EXPECT_GT(5, order["const_node3"]);
EXPECT_GT(5, order["const_node4"]);
}
void TestCommonAncestor() {
GraphDef graph_def;
NodeDef* add_node1 = graph_def.add_node();
add_node1->set_name("add_node1");
add_node1->set_op("Add");
add_node1->add_input("add_node2");
add_node1->add_input("add_node3");
NodeDef* add_node2 = graph_def.add_node();
add_node2->set_name("add_node2");
add_node2->set_op("Add");
add_node2->add_input("const_node1");
add_node2->add_input("const_node2");
NodeDef* add_node3 = graph_def.add_node();
add_node3->set_name("add_node3");
add_node3->set_op("Add");
add_node3->add_input("const_node1");
add_node3->add_input("const_node3");
NodeDef* const_node1 = graph_def.add_node();
const_node1->set_name("const_node1");
const_node1->set_op("Const");
NodeDef* const_node2 = graph_def.add_node();
const_node2->set_name("const_node2");
const_node2->set_op("Const");
NodeDef* const_node3 = graph_def.add_node();
const_node3->set_name("const_node3");
const_node3->set_op("Const");
GraphDef result;
TF_ASSERT_OK(SortByExecutionOrder(graph_def, &result));
std::map<std::string, int> order;
GetOrder(result, &order);
EXPECT_EQ(5, order["add_node1"]);
EXPECT_GT(5, order["add_node2"]);
EXPECT_GT(5, order["add_node3"]);
EXPECT_GT(4, order["const_node2"]);
EXPECT_GT(4, order["const_node3"]);
EXPECT_GT(3, order["const_node1"]);
}
};
TEST_F(SortByExecutionOrderTest, TestSimpleAdd) { TestSimpleAdd(); }
TEST_F(SortByExecutionOrderTest, TestSimpleLinear) { TestSimpleLinear(); }
TEST_F(SortByExecutionOrderTest, TestSimpleTree) { TestSimpleTree(); }
TEST_F(SortByExecutionOrderTest, TestCommonAncestor) { TestCommonAncestor(); }
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/sparsify_gather.cc
+629
View File
@@ -0,0 +1,629 @@
/* Copyright 2016 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 <cmath>
#include <memory>
#include <unordered_map>
#include "tensorflow/c/checkpoint_reader.h"
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/graph/node_builder.h"
#include "tensorflow/core/graph/subgraph.h"
#include "tensorflow/core/lib/strings/str_util.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/core/util/tensor_bundle/tensor_bundle.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
using str_util::Split;
using str_util::StringReplace;
using strings::StrCat;
namespace graph_transforms {
// Sparsify Tensor of shape [N, 1]. Return the indices and values vectors for
// non-zero tensor content.
absl::Status SparsifyWeights(const Tensor& tensor, Tensor* indices_tensor,
Tensor* values_tensor) {
if (tensor.dims() != 2 || tensor.dim_size(1) != 1) {
return absl::FailedPreconditionError(
absl::StrCat("Transform only applicable to subgraph with 'Const' with "
"tensor of shape [N, 1]. But instead get shape ",
tensor.shape().DebugString(), "."));
}
auto flat = tensor.flat<float>();
std::vector<int64_t> indices;
std::vector<float> values;
for (int64_t i = 0; i < flat.size(); i++) {
float val = flat(i);
if (std::abs(val) >= 1.0e-5) {
indices.push_back(i);
values.push_back(val);
}
}
// During model initialization, InitializeTableOp makes use of
// KeyValueTensorIterator, which does not accept empty keys or values.
// Consequently, adding a dummy pair of indices and values as a walkaround.
if (indices.empty() || values.empty()) {
indices.push_back(0);
values.push_back(0);
}
*indices_tensor = Tensor(DataTypeToEnum<int64_t>::value,
{static_cast<int64_t>(indices.size())});
std::copy_n(indices.begin(), indices.size(),
indices_tensor->flat<int64_t>().data());
*values_tensor = Tensor(DataTypeToEnum<float>::value,
{static_cast<int64_t>(values.size())});
std::copy_n(values.begin(), values.size(),
values_tensor->flat<float>().data());
return absl::OkStatus();
}
void CreateConstNode(const Tensor& tensor, const std::string& name,
NodeDef* node_def) {
node_def->set_op("Const");
node_def->set_name(name);
SetNodeTensorAttr<float>("value", tensor, node_def);
}
std::string GetMonolithicTensorKey(const std::string& tensor_slice_name) {
std::vector<std::string> names = Split(tensor_slice_name, "/");
if (absl::StartsWith(names[names.size() - 1], "part_")) {
CHECK_GE(names.size(), 2);
names.pop_back();
}
return absl::StrJoin(names, "/");
}
absl::Status ObtainTensorSlice(const GraphDef& input_graph_def,
const std::string& target_name,
std::string* shape_slice_string) {
std::string restore_node_name;
for (const auto& node : input_graph_def.node()) {
std::vector<std::string> node_name_parts = Split(node.name(), "/");
if (node_name_parts.size() == 2 &&
absl::StartsWith(node_name_parts[0], "save") &&
absl::StartsWith(node_name_parts[1], "Assign") &&
node.input(0) == target_name) {
restore_node_name = node.input(1);
break;
}
}
std::vector<std::string> restore_node_parts = Split(restore_node_name, ":");
CHECK_LE(restore_node_parts.size(), 2);
std::string tensor_names_node;
std::string shape_and_slices_node;
for (const auto& node : input_graph_def.node()) {
if ((node.name() == restore_node_parts[0]) && (node.op() == "RestoreV2")) {
tensor_names_node = node.input(1);
shape_and_slices_node = node.input(2);
break;
}
}
int offset = -1;
for (const auto& node : input_graph_def.node()) {
if (node.name() == tensor_names_node) {
Tensor tensor_names_tensor;
TF_RETURN_IF_ERROR(GetNodeAttr(node, "value", &tensor_names_tensor));
const auto& tensor_names_value = tensor_names_tensor.flat<tstring>();
for (int i = 0; i < tensor_names_value.size(); i++) {
if (tensor_names_value(i) == GetMonolithicTensorKey(target_name)) {
offset = i;
break;
}
}
}
}
if (offset == -1) {
return absl::InternalError(absl::StrCat(
"Unable to find RestoreV2 entry for variable: ", target_name));
}
for (const auto& node : input_graph_def.node()) {
if (node.name() == shape_and_slices_node) {
Tensor shape_and_slices_tensor;
TF_RETURN_IF_ERROR(GetNodeAttr(node, "value", &shape_and_slices_tensor));
const auto& shape_and_slices_value =
shape_and_slices_tensor.flat<tstring>();
*shape_slice_string = shape_and_slices_value(offset);
return absl::OkStatus();
}
}
return absl::InternalError(
absl::StrCat("Unable to find slice for variable: ", target_name));
}
absl::Status ReadTensorFromCheckpoint(
const std::string& tensor_name,
const std::unique_ptr<BundleReader>& ckpt_reader,
const std::string& shape_and_slice, Tensor* tensor) {
if (ckpt_reader) {
TensorShape parsed_full_shape;
TensorSlice parsed_slice;
TensorShape parsed_slice_shape;
bool get_slice = false;
if (!shape_and_slice.empty()) {
TF_RETURN_IF_ERROR(
checkpoint::ParseShapeAndSlice(shape_and_slice, &parsed_full_shape,
&parsed_slice, &parsed_slice_shape));
get_slice = (parsed_full_shape != parsed_slice_shape);
}
if (get_slice) {
TF_RETURN_IF_ERROR(ckpt_reader->LookupSlice(
GetMonolithicTensorKey(tensor_name), parsed_slice, tensor));
} else {
TF_RETURN_IF_ERROR(
ckpt_reader->Lookup(GetMonolithicTensorKey(tensor_name), tensor));
}
return absl::OkStatus();
}
return absl::InternalError("Checkpoint reader was not initialized. ");
}
absl::Status InitializeCheckpointReader(
const TransformFuncContext& context,
std::unique_ptr<BundleReader>* ckpt_reader) {
if (context.params.count("input_checkpoint")) {
const std::string input_checkpoint =
context.params.at("input_checkpoint")[0];
*ckpt_reader =
std::make_unique<BundleReader>(Env::Default(), input_checkpoint);
TF_RETURN_IF_ERROR((*ckpt_reader)->status());
}
return absl::OkStatus();
}
absl::Status ObtainVariableInfo(
const GraphDef& input_graph_def,
std::unique_ptr<std::unordered_map<std::string, std::string>>*
shapes_and_slices) {
*shapes_and_slices =
std::make_unique<std::unordered_map<std::string, std::string>>();
for (const auto& node : input_graph_def.node()) {
if ((node.op() == "Variable") || (node.op() == "VariableV2")) {
std::string s;
TF_RETURN_IF_ERROR(ObtainTensorSlice(input_graph_def, node.name(), &s));
(**shapes_and_slices)[node.name()] = s;
}
}
return absl::OkStatus();
}
absl::Status RemoveInputAtIndex(NodeDef* n, int index) {
for (int i = index; i < n->input_size() - 1; i++) {
n->mutable_input()->SwapElements(i, i + 1);
}
n->mutable_input()->RemoveLast();
return absl::OkStatus();
}
absl::Status RemoveNodeAtIndex(GraphDef* g, int index) {
for (int i = index; i < g->node_size() - 1; i++) {
g->mutable_node()->SwapElements(i, i + 1);
}
g->mutable_node()->RemoveLast();
return absl::OkStatus();
}
absl::Status SparsifyGatherInternal(
const GraphDef& input_graph_def,
const std::unique_ptr<std::unordered_map<std::string, std::string>>&
shapes_and_slices,
const TransformFuncContext& context, const OpTypePattern& pattern,
const std::unique_ptr<BundleReader>& ckpt_reader,
GraphDef* output_graph_def) {
std::string group_init_node = "group_deps";
if (context.params.count("group_init_node")) {
group_init_node = context.params.at("group_init_node")[0];
}
GraphDef current_graph_def = input_graph_def;
bool any_match_found = false;
// Populate references.
std::unordered_map<std::string, int> refs;
for (const auto& node : current_graph_def.node()) {
for (const auto& input : node.input()) {
auto parsed_input = StringReplace(input, "^", "", true);
refs[parsed_input] += 1;
}
}
// The subgraphs may have overlapping components, therefore GraphMatcher
// doesn't return all subgraphs in one round -- this has to be multi-round
// update.
do {
any_match_found = false;
GraphDef replaced_graph_def = current_graph_def;
std::vector<std::string> init_table_node_names;
std::vector<std::string> removed_node_names;
TF_RETURN_IF_ERROR(ReplaceMatchingOpTypes(
current_graph_def, pattern,
[&ckpt_reader, &any_match_found, &init_table_node_names,
&shapes_and_slices, &removed_node_names, &refs](
const NodeMatch& match, const std::set<std::string>& input_nodes,
const std::set<std::string>& output_nodes,
std::vector<NodeDef>* new_nodes) {
any_match_found = true;
// The captured subgraph should be of the following pattern:
// Const --> Identity --> Gather --> ...
// ^
// |
// (ids)
//
// After transform, it becomes:
// --> NoOp(group_deps)
// |
// Const --> InitializeTable --> HashTable
// ^ |
// | |
// Const ------------- |
// v
// (ids) ---> LookupTableFind <--- Const(default)
// |
// v
// ...
// clang-format off
// For each subgraph, do the following
// 1. Sparsify the `Const`, creating two `Const`, for hashtable
// key/val.
// 2. Create a `InitializeTable` op connecting to the above 2 `Const`.
// 3. Create a `HashTable` op connecting to `InitializeTable` op.
// 4. Replace the `Gather` with a `LookupTableFind` op.
// 5. Connect the `LookupTableFind` with
// a. `HashTable`
// b. `Gather`'s ids input
// c. a `default_val` arg, valued at 0
// clang-format on
const NodeDef& gather_node = match.node;
// GatherV2 adds an "axis" parameter. sparsify_gather only supports
// axis 0 gathers.
if (gather_node.op() == "GatherV2") {
// Per the OpTypePattern, the 3rd input to Gather must be a Const.
const NodeDef& axis_node = match.inputs[2].node;
Tensor axis_t;
TF_RETURN_IF_ERROR(GetNodeAttr(axis_node, "value", &axis_t));
int64_t axis = 0;
if (axis_t.dtype() == DT_INT32) {
axis = axis_t.scalar<int32_t>()();
} else if (axis_t.dtype() == DT_INT64) {
axis = axis_t.scalar<int64_t>()();
} else {
return absl::FailedPreconditionError(
"Gather axis was not int32 or int64.");
}
if (axis != 0) {
return absl::FailedPreconditionError(absl::StrCat(
"Transform only applicable to subgraph with GatherV2 over "
"axis 0. Found axis ",
axis, "."));
}
}
const NodeDef& weights_node = match.inputs[0].inputs[0].node;
DataType data_type;
TF_RETURN_IF_ERROR(GetNodeAttr(weights_node, "dtype", &data_type));
if (data_type != DT_FLOAT) {
return absl::FailedPreconditionError(absl::StrCat(
"Transform only applicable to subgraph with 'Const',"
"'Variable', or 'VariableV2' of dtype "
"'DT_FLOAT'. Found '" +
weights_node.op() + "' with name '",
weights_node.name(), "' and dtype '", data_type, "'."));
}
Tensor weight;
if (weights_node.op() == "Const") {
weight = GetNodeTensorAttr(weights_node, "value");
} else {
TF_RETURN_IF_ERROR(ReadTensorFromCheckpoint(
weights_node.name(), ckpt_reader,
(*shapes_and_slices)[weights_node.name()], &weight));
}
// Add both weight and identity node names.
removed_node_names.push_back(weights_node.name());
removed_node_names.push_back(match.inputs[0].node.name());
for (auto input_node : match.inputs[0].node.input()) {
auto parsed_input = StringReplace(input_node, "^", "", true);
refs[parsed_input]--;
}
Tensor indices_tensor;
Tensor values_tensor;
TF_RETURN_IF_ERROR(
SparsifyWeights(weight, &indices_tensor, &values_tensor));
// indices and values of sparsified `Const`
DataType key_dtype = DT_INT64;
NodeDef indices_node;
CreateConstNode(indices_tensor,
absl::StrCat(weights_node.name(), "/indices"),
&indices_node);
SetNodeAttr("dtype", key_dtype, &indices_node);
NodeDef values_node;
CreateConstNode(values_tensor,
absl::StrCat(weights_node.name(), "/values"),
&values_node);
SetNodeAttr("dtype", data_type, &values_node);
// HashTable node
NodeDef hashtable_node;
hashtable_node.set_op("HashTable");
hashtable_node.set_name(
absl::StrCat(weights_node.name(), "/HashTable"));
SetNodeAttr("key_dtype", key_dtype, &hashtable_node);
SetNodeAttr("value_dtype", data_type, &hashtable_node);
// InitializeTable node
NodeDef init_table_node;
init_table_node.set_op("InitializeTable");
init_table_node.set_name(
absl::StrCat(weights_node.name(), "/InitializeTable"));
SetNodeAttr("Tkey", key_dtype, &init_table_node);
SetNodeAttr("Tval", data_type, &init_table_node);
init_table_node_names.push_back(init_table_node.name());
// LookupTableFind node
NodeDef lookup_node;
lookup_node.set_op("LookupTableFind");
lookup_node.set_name(
absl::StrCat(gather_node.name(), "/LookupTableFind"));
SetNodeAttr("Tin", key_dtype, &lookup_node);
SetNodeAttr("Tout", data_type, &lookup_node);
// Default return value of hashtable lookup
Tensor zero_tensor(data_type, TensorShape({}));
zero_tensor.flat<float>()(0) = 0.0;
NodeDef default_value_node;
CreateConstNode(zero_tensor,
absl::StrCat(gather_node.name(), "/Const"),
&default_value_node);
SetNodeAttr("dtype", data_type, &default_value_node);
// ExpandDims argument
Tensor dim_idx(DT_INT32, TensorShape({}));
dim_idx.flat<int32_t>()(0) = -1;
NodeDef dim_idx_node;
dim_idx_node.set_op("Const");
dim_idx_node.set_name(
absl::StrCat(gather_node.name(), "/ExpandDims/Const"));
SetNodeAttr("value", dim_idx, &dim_idx_node);
SetNodeAttr("dtype", DT_INT32, &dim_idx_node);
// ExpandDims node
NodeDef expand_dims_node;
expand_dims_node.set_op("ExpandDims");
// Reuse gather_node's name so not to change dependent's inputs
expand_dims_node.set_name(gather_node.name());
SetNodeAttr("T", data_type, &expand_dims_node);
// Connect nodes
AddNodeInput(hashtable_node.name(), &init_table_node);
refs[hashtable_node.name()]++;
AddNodeInput(indices_node.name(), &init_table_node);
refs[indices_node.name()]++;
AddNodeInput(values_node.name(), &init_table_node);
refs[values_node.name()]++;
AddNodeInput(hashtable_node.name(), &lookup_node);
refs[hashtable_node.name()]++;
AddNodeInput(gather_node.input(1), &lookup_node);
refs[gather_node.input(1)]++;
AddNodeInput(default_value_node.name(), &lookup_node);
refs[default_value_node.name()]++;
AddNodeInput(lookup_node.name(), &expand_dims_node);
refs[lookup_node.name()]++;
AddNodeInput(dim_idx_node.name(), &expand_dims_node);
refs[dim_idx_node.name()]++;
// Copy 'ids' input of original 'Gather'
new_nodes->push_back(match.inputs[1].node);
new_nodes->push_back(indices_node);
new_nodes->push_back(values_node);
new_nodes->push_back(hashtable_node);
new_nodes->push_back(init_table_node);
new_nodes->push_back(lookup_node);
new_nodes->push_back(default_value_node);
new_nodes->push_back(dim_idx_node);
new_nodes->push_back(expand_dims_node);
return absl::OkStatus();
},
{true}, &replaced_graph_def));
NodeDef* init_op = nullptr;
for (int i = 0; i < replaced_graph_def.node_size(); i++) {
if (replaced_graph_def.node(i).name() == group_init_node &&
replaced_graph_def.node(i).op() == "NoOp") {
init_op = replaced_graph_def.mutable_node(i);
break;
}
}
if (!init_op) {
// Init node
init_op = replaced_graph_def.mutable_node()->Add();
init_op->set_op("NoOp");
init_op->set_name(group_init_node);
}
for (const std::string& name : init_table_node_names) {
// Add control dependence from init_table_node to group_deps_node
AddNodeInput(absl::StrCat("^", name), init_op);
refs[name]++;
}
// Erase inputs and outputs as they are not considered for deletion.
for (const auto& output : context.output_names) {
refs.erase(output);
}
for (const auto& input : context.input_names) {
refs.erase(input);
}
// Add nodes with a reference count of 0 for deletion.
for (const auto& entry : refs) {
if (entry.second == 0) {
removed_node_names.push_back(entry.first);
}
}
while (!removed_node_names.empty()) {
auto name = removed_node_names.back();
removed_node_names.pop_back();
int i = 0;
while (i < replaced_graph_def.node_size()) {
// Revisit this to see if we can safely remove RestoreV2 nodes.
if ((replaced_graph_def.node(i).name() == name) &&
(replaced_graph_def.node(i).op() != "RestoreV2")) {
for (const auto& input : replaced_graph_def.node(i).input()) {
auto parsed_input = StringReplace(input, "^", "", true);
refs[parsed_input] -= 1;
if (refs[parsed_input] == 0) {
removed_node_names.push_back(parsed_input);
}
}
TF_RETURN_IF_ERROR(RemoveNodeAtIndex(&replaced_graph_def, i));
continue;
}
int j = 0;
bool deleted_inputs = false;
while (j < replaced_graph_def.node(i).input_size()) {
if (replaced_graph_def.node(i).input(j) == name ||
replaced_graph_def.node(i).input(j) == ("^" + name)) {
TF_RETURN_IF_ERROR(
RemoveInputAtIndex(replaced_graph_def.mutable_node(i), j));
deleted_inputs = true;
continue;
}
j++;
}
if (deleted_inputs) {
if (replaced_graph_def.node(i).op() == "ConcatV2") {
if (replaced_graph_def.node(i).input_size() > 2) {
SetNodeAttr("N", replaced_graph_def.node(i).input_size() - 1,
replaced_graph_def.mutable_node(i));
} else if (replaced_graph_def.node(i).input_size() == 2) {
if (refs[replaced_graph_def.node(i).input(1)] != 1) {
return absl::InternalError(
"Expect axis tensor of ConcatV2 node to only be referenced "
"once.");
}
refs[replaced_graph_def.node(i).input(1)] -= 1;
removed_node_names.push_back(replaced_graph_def.node(i).input(1));
replaced_graph_def.mutable_node(i)->mutable_input()->RemoveLast();
replaced_graph_def.mutable_node(i)->mutable_attr()->erase("N");
replaced_graph_def.mutable_node(i)->set_op("Identity");
} else {
return absl::InternalError(
"ConcatV2 should have at least two elements");
}
}
if ((replaced_graph_def.node(i).op() == "Assign" ||
replaced_graph_def.node(i).op() == "Reshape" ||
replaced_graph_def.node(i).op() == "Equal" ||
replaced_graph_def.node(i).op() == "Mean" ||
replaced_graph_def.node(i).op() == "ScalarSummary") &&
replaced_graph_def.node(i).input_size() == 1) {
removed_node_names.push_back(replaced_graph_def.node(i).name());
}
if (!replaced_graph_def.node(i).input_size()) {
removed_node_names.push_back(replaced_graph_def.node(i).name());
}
}
i++;
}
}
current_graph_def = replaced_graph_def;
} while (any_match_found);
*output_graph_def = current_graph_def;
return absl::OkStatus();
}
absl::Status SparsifyGather(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def) {
// clang-format off
const OpTypePattern gather_pattern =
{"Gather",
{
{"Identity",
{
{"Const|Variable|VariableV2"}
}
},
{"*"},
}
};
const OpTypePattern gather_v2_pattern =
{"GatherV2",
{
{"Identity",
{
{"Const|Variable|VariableV2"}
}
},
{"*"},
// GatherV2's axis must be constant.
{"Const"},
}
};
// clang-format on
GraphDef cleaned_input_graph_def;
RemoveAttributes(input_graph_def, {"_output_shapes"},
&cleaned_input_graph_def);
GraphDef temp_output;
std::unique_ptr<BundleReader> ckpt_reader;
TF_RETURN_IF_ERROR(InitializeCheckpointReader(context, &ckpt_reader));
std::unique_ptr<std::unordered_map<std::string, std::string>>
shapes_and_slices;
TF_RETURN_IF_ERROR(
ObtainVariableInfo(cleaned_input_graph_def, &shapes_and_slices));
TF_RETURN_IF_ERROR(SparsifyGatherInternal(
cleaned_input_graph_def, shapes_and_slices, context, gather_pattern,
ckpt_reader, &temp_output));
TF_RETURN_IF_ERROR(SparsifyGatherInternal(temp_output, shapes_and_slices,
context, gather_v2_pattern,
ckpt_reader, output_graph_def));
return absl::OkStatus();
}
REGISTER_GRAPH_TRANSFORM("sparsify_gather", SparsifyGather);
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/sparsify_gather_test.cc
+643
View File
@@ -0,0 +1,643 @@
/* Copyright 2015 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/cc/ops/const_op.h"
#include "tensorflow/cc/ops/sendrecv_ops.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/lib/io/path.h"
#include "tensorflow/core/lib/strings/strcat.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/core/util/tensor_bundle/tensor_bundle.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
// Declarations so we don't need a public header.
absl::Status SparsifyGather(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def);
absl::Status ReadTensorFromCheckpoint(
const std::string& tensor_name,
const std::unique_ptr<BundleReader>& ckpt_reader,
const std::string& shape_and_slice, Tensor* tensor);
class SparsifyGatherTest : public ::testing::Test {
protected:
NodeDef* CreateNode(const absl::string_view name, const absl::string_view op,
const std::vector<NodeDef*>& inputs, GraphDef* graph_def,
bool control_dep = false) {
NodeDef* node_def = graph_def->add_node();
node_def->set_name(name);
node_def->set_op(op);
if (!control_dep) {
std::for_each(inputs.begin(), inputs.end(), [&node_def](NodeDef* input) {
node_def->add_input(input->name());
});
} else {
std::for_each(inputs.begin(), inputs.end(), [&node_def](NodeDef* input) {
node_def->add_input(absl::StrCat("^", input->name()));
});
}
return node_def;
}
void MakeGather(absl::string_view name, bool gather_v2, NodeDef* params,
NodeDef* indices, GraphDef* graph_def) {
if (gather_v2) {
NodeDef* axis_node =
CreateNode(absl::StrCat(name, "_axis"), "Const", {}, graph_def);
Tensor axis_t(DT_INT32, TensorShape({}));
axis_t.scalar<int32_t>()() = 0;
SetNodeTensorAttr<int32_t>("value", axis_t, axis_node);
CreateNode(name, "GatherV2", {params, indices, axis_node}, graph_def);
} else {
CreateNode(name, "Gather", {params, indices}, graph_def);
}
}
void TestSinglePartition(bool gather_v2, bool include_shared_init,
bool test_variable, bool test_kept_concat,
const std::string& shared_init_name = "group_deps") {
GraphDef graph_def;
const auto checkpoint_path =
io::JoinPath(testing::TmpDir(), "checkpoint_single");
// Build the graph.
NodeDef* input_node = CreateNode("ids", "Const", {}, &graph_def);
NodeDef* w_node;
NodeDef* zeros_const;
NodeDef* zeros_shape;
NodeDef* zeros_node;
NodeDef* assign_node;
Tensor weights(DT_FLOAT, TensorShape({4, 1}));
test::FillValues<float>(&weights, {0.2, 0.000001, 1.2, 0.001});
if (!test_variable) {
w_node = CreateNode("w/part_1", "Const", {}, &graph_def);
SetNodeTensorAttr<float>("value", weights, w_node);
} else {
w_node = CreateNode("w/part_1", "VariableV2", {}, &graph_def);
zeros_shape = CreateNode("w/part_1/Initializer/zeros/shape_as_tensor",
"Const", {}, &graph_def);
zeros_const = CreateNode("w/part_1/Initializer/zeros/Const", "Const", {},
&graph_def);
zeros_node = CreateNode("w/part_1/Initializer/zeros", "Fill",
{zeros_shape, zeros_const}, &graph_def);
assign_node = CreateNode("w/part_1/Assign", "Assign",
{w_node, zeros_node}, &graph_def);
NodeDef* save_const_node =
CreateNode("save/Const", "Const", {}, &graph_def);
Tensor tensor_names_values(DT_STRING, TensorShape({1}));
test::FillValues<tstring>(&tensor_names_values, {"w"});
NodeDef* tensor_names_node =
CreateNode("save/RestoreV2/tensor_names", "Const", {}, &graph_def);
SetNodeTensorAttr<std::string>("value", tensor_names_values,
tensor_names_node);
NodeDef* tensor_shapes_slices_node = CreateNode(
"save/RestoreV2/shape_and_slices", "Const", {}, &graph_def);
Tensor shapes_slices_val(DT_STRING, TensorShape({1}));
shapes_slices_val.flat<tstring>()(0) = "4 1 0,4:0,1";
SetNodeTensorAttr<std::string>("value", shapes_slices_val,
tensor_shapes_slices_node);
NodeDef* restore_node = CreateNode(
"save/RestoreV2", "RestoreV2",
{save_const_node, tensor_names_node, tensor_shapes_slices_node},
&graph_def);
CreateNode("save/Assign", "Assign", {w_node, restore_node}, &graph_def);
BundleWriter writer(Env::Default(), checkpoint_path);
TF_ASSERT_OK(writer.Add("w", weights));
TF_ASSERT_OK(writer.Finish());
}
SetNodeAttr("dtype", DT_FLOAT, w_node);
NodeDef* identity_node =
CreateNode("w/read", "Identity", {w_node}, &graph_def);
MakeGather("gather", gather_v2, identity_node, input_node, &graph_def);
if (include_shared_init) {
if (!test_variable) {
CreateNode(shared_init_name, "NoOp", {}, &graph_def);
} else {
CreateNode(shared_init_name, "NoOp", {assign_node}, &graph_def, true);
}
}
NodeDef* concat_axis_node =
CreateNode("linear/concat/axis", "Const", {}, &graph_def);
NodeDef* concat_input_node =
CreateNode("concat/input/node", "Const", {}, &graph_def);
NodeDef* concat_node = nullptr;
if (!test_kept_concat) {
concat_node = CreateNode(
"concat/node", "ConcatV2",
{identity_node, concat_input_node, concat_axis_node}, &graph_def);
SetNodeAttr("N", 2, concat_node);
} else {
NodeDef* concat_input_node_2 =
CreateNode("concat/input/node_2", "Const", {}, &graph_def);
concat_node = CreateNode("concat/node", "ConcatV2",
{identity_node, concat_input_node,
concat_input_node_2, concat_axis_node},
&graph_def);
SetNodeAttr("N", 3, concat_node);
}
// Run the op.
GraphDef result;
TransformFuncContext context;
context.input_names = {"ids"};
context.output_names = {"gather"};
if (test_variable) {
context.params["input_checkpoint"] = {checkpoint_path};
}
if (shared_init_name != "group_deps") {
context.params["group_init_node"] = {shared_init_name};
}
TF_ASSERT_OK(SparsifyGather(graph_def, context, &result));
// Validation begins.
std::map<std::string, const NodeDef*> node_lookup;
MapNamesToNodes(result, &node_lookup);
// Check nodes.
EXPECT_EQ(0,
node_lookup.count("w/part_1/Initializer/zeros/shape_as_tensor"));
EXPECT_EQ(0, node_lookup.count("w/part_1/Initializer/zeros/Const"));
EXPECT_EQ(0, node_lookup.count("w/part_1/Initializer/zeros"));
EXPECT_EQ(0, node_lookup.count("w/part_1/Assign"));
EXPECT_EQ(1, node_lookup.count("ids"));
EXPECT_EQ("Const", node_lookup.at("ids")->op());
EXPECT_EQ(1, node_lookup.count("concat/node"));
if (!test_kept_concat) {
EXPECT_EQ(0, node_lookup.count("linear/concat/axis"));
EXPECT_EQ("Identity", node_lookup.at("concat/node")->op());
EXPECT_EQ(1, node_lookup.at("concat/node")->input_size());
EXPECT_EQ("concat/input/node", node_lookup.at("concat/node")->input(0));
} else {
EXPECT_EQ(1, node_lookup.count("linear/concat/axis"));
EXPECT_EQ("ConcatV2", node_lookup.at("concat/node")->op());
EXPECT_EQ(3, node_lookup.at("concat/node")->input_size());
EXPECT_EQ("concat/input/node", node_lookup.at("concat/node")->input(0));
EXPECT_EQ("concat/input/node_2", node_lookup.at("concat/node")->input(1));
EXPECT_EQ("linear/concat/axis", node_lookup.at("concat/node")->input(2));
EXPECT_EQ(2, node_lookup.at("concat/node")->attr().at("N").i());
}
EXPECT_EQ(1, node_lookup.count("w/part_1/indices"));
EXPECT_EQ("Const", node_lookup.at("w/part_1/indices")->op());
Tensor expected_indices_tensor(DT_INT64, TensorShape({3}));
test::FillValues<int64_t>(&expected_indices_tensor, {0, 2, 3});
test::ExpectTensorEqual<int64_t>(
expected_indices_tensor,
GetNodeTensorAttr(*(node_lookup.at("w/part_1/indices")), "value"));
EXPECT_EQ(1, node_lookup.count("w/part_1/values"));
EXPECT_EQ("Const", node_lookup.at("w/part_1/values")->op());
Tensor expected_values_tensor(DT_FLOAT, TensorShape({3}));
test::FillValues<float>(&expected_values_tensor, {0.2, 1.2, 0.001});
test::ExpectTensorNear<float>(
expected_values_tensor,
GetNodeTensorAttr(*(node_lookup.at("w/part_1/values")), "value"), 1e-5);
EXPECT_EQ(1, node_lookup.count("w/part_1/HashTable"));
EXPECT_EQ("HashTable", node_lookup.at("w/part_1/HashTable")->op());
EXPECT_EQ(1, node_lookup.count("w/part_1/InitializeTable"));
EXPECT_EQ("InitializeTable",
node_lookup.at("w/part_1/InitializeTable")->op());
// Nodes in "gather" scope.
EXPECT_EQ(1, node_lookup.count("gather/LookupTableFind"));
EXPECT_EQ("LookupTableFind",
node_lookup.at("gather/LookupTableFind")->op());
EXPECT_EQ(1, node_lookup.count("gather/Const"));
EXPECT_EQ("Const", node_lookup.at("gather/Const")->op());
Tensor expected_gather_default_tensor(DT_FLOAT, TensorShape({}));
test::FillValues<float>(&expected_gather_default_tensor, {0.0});
test::ExpectTensorNear<float>(
expected_gather_default_tensor,
GetNodeTensorAttr(*(node_lookup.at("gather/Const")), "value"), 1e-5);
EXPECT_EQ(1, node_lookup.count("gather/ExpandDims/Const"));
EXPECT_EQ("Const", node_lookup.at("gather/ExpandDims/Const")->op());
Tensor expected_expand_dims_tensor(DT_INT32, TensorShape({}));
test::FillValues<int32_t>(&expected_expand_dims_tensor, {-1});
test::ExpectTensorEqual<int32_t>(
expected_expand_dims_tensor,
GetNodeTensorAttr(*(node_lookup.at("gather/ExpandDims/Const")),
"value"));
EXPECT_EQ(1, node_lookup.count("gather"));
EXPECT_EQ("ExpandDims", node_lookup.at("gather")->op());
EXPECT_EQ(1, node_lookup.count(shared_init_name));
EXPECT_EQ("NoOp", node_lookup.at(shared_init_name)->op());
// Check connections
EXPECT_EQ("w/part_1/HashTable",
node_lookup.at("w/part_1/InitializeTable")->input(0));
EXPECT_EQ("w/part_1/indices",
node_lookup.at("w/part_1/InitializeTable")->input(1));
EXPECT_EQ("w/part_1/values",
node_lookup.at("w/part_1/InitializeTable")->input(2));
EXPECT_EQ("w/part_1/HashTable",
node_lookup.at("gather/LookupTableFind")->input(0));
EXPECT_EQ("ids", node_lookup.at("gather/LookupTableFind")->input(1));
EXPECT_EQ("gather/Const",
node_lookup.at("gather/LookupTableFind")->input(2));
EXPECT_EQ("gather/LookupTableFind", node_lookup.at("gather")->input(0));
// Check control dependency.
EXPECT_NE(std::find(node_lookup.at(shared_init_name)->input().begin(),
node_lookup.at(shared_init_name)->input().end(),
"^w/part_1/InitializeTable"),
node_lookup.at(shared_init_name)->input().end());
EXPECT_EQ(1, node_lookup.at(shared_init_name)->input().size());
}
void TestMultiPartition(bool gather_v2, bool include_shared_init,
bool test_variable,
const std::string& shared_init_name = "group_deps") {
// The 'ids' node is served input for two 'Gather's.
GraphDef graph_def;
const auto checkpoint_path =
io::JoinPath(testing::TmpDir(), "checkpoint_multiple");
// Build Graph:
// Shared input node
NodeDef* input_node = CreateNode("ids", "Const", {}, &graph_def);
// Two partitions
NodeDef* w_node1;
NodeDef* w_node2;
NodeDef* zeros_const1;
NodeDef* zeros_shape1;
NodeDef* zeros_node1;
NodeDef* zeros_const2;
NodeDef* zeros_shape2;
NodeDef* zeros_node2;
NodeDef* assign_node1;
NodeDef* assign_node2;
Tensor weights(DT_FLOAT, TensorShape({4, 1}));
test::FillValues<float>(&weights, {0.2, 0.000001, 1.2, 0.001});
if (!test_variable) {
w_node1 = CreateNode("w1/part_1", "Const", {}, &graph_def);
w_node2 = CreateNode("w2/part_1", "Const", {}, &graph_def);
SetNodeTensorAttr<float>("value", weights, w_node1);
SetNodeTensorAttr<float>("value", weights, w_node2);
} else {
NodeDef* save_const_node =
CreateNode("save/Const", "Const", {}, &graph_def);
NodeDef* tensor_names_node =
CreateNode("save/RestoreV2/tensor_names", "Const", {}, &graph_def);
Tensor tensor_names_values(DT_STRING, TensorShape({2}));
test::FillValues<tstring>(&tensor_names_values, {"w1", "w2"});
SetNodeTensorAttr<std::string>("value", tensor_names_values,
tensor_names_node);
NodeDef* tensor_shapes_slices_node = CreateNode(
"save/RestoreV2/shape_and_slices", "Const", {}, &graph_def);
Tensor shapes_slices_val(DT_STRING, TensorShape({2}));
shapes_slices_val.flat<tstring>()(0) = "4 1 0,4:0,1";
shapes_slices_val.flat<tstring>()(1) = "4 1 0,4:0,1";
SetNodeTensorAttr<std::string>("value", shapes_slices_val,
tensor_shapes_slices_node);
NodeDef* restore_node = CreateNode(
"save/RestoreV2", "RestoreV2",
{save_const_node, tensor_names_node, tensor_shapes_slices_node},
&graph_def);
w_node1 = CreateNode("w1/part_1", "VariableV2", {}, &graph_def);
zeros_shape1 = CreateNode("w1/part_1/Initializer/zeros/shape_as_tensor",
"Const", {}, &graph_def);
zeros_const1 = CreateNode("w1/part_1/Initializer/zeros/Const", "Const",
{}, &graph_def);
zeros_node1 = CreateNode("w1/part_1/Initializer/zeros", "Fill",
{zeros_shape1, zeros_const1}, &graph_def);
assign_node1 = CreateNode("w1/part_1/Assign", "Assign",
{w_node1, zeros_node1}, &graph_def);
CreateNode("save/Assign", "Assign", {w_node1, restore_node}, &graph_def);
w_node2 = CreateNode("w2/part_1", "VariableV2", {}, &graph_def);
zeros_shape2 = CreateNode("w2/part_1/Initializer/zeros/shape_as_tensor",
"Const", {}, &graph_def);
zeros_const2 = CreateNode("w2/part_1/Initializer/zeros/Const", "Const",
{}, &graph_def);
zeros_node2 = CreateNode("w2/part_1/Initializer/zeros", "Fill",
{zeros_shape2, zeros_const2}, &graph_def);
assign_node2 = CreateNode("w2/part_1/Assign", "Assign",
{w_node2, zeros_node2}, &graph_def);
CreateNode("save/Assign_1", "Assign", {w_node2, restore_node},
&graph_def);
BundleWriter writer(Env::Default(), checkpoint_path);
TF_ASSERT_OK(writer.Add("w1", weights));
TF_ASSERT_OK(writer.Add("w2", weights));
TF_ASSERT_OK(writer.Finish());
}
SetNodeAttr("dtype", DT_FLOAT, w_node1);
SetNodeAttr("dtype", DT_FLOAT, w_node2);
NodeDef* identity_node1 =
CreateNode("w1/part_1/read", "Identity", {w_node1}, &graph_def);
NodeDef* identity_node2 =
CreateNode("w2/part_1/read", "Identity", {w_node2}, &graph_def);
MakeGather("gather1", gather_v2, identity_node1, input_node, &graph_def);
MakeGather("gather2", gather_v2, identity_node2, input_node, &graph_def);
NodeDef* concat_axis_node =
CreateNode("linear/concat/axis", "Const", {}, &graph_def);
NodeDef* concat_node = CreateNode(
"concat/node", "ConcatV2",
{identity_node1, identity_node2, concat_axis_node}, &graph_def);
SetNodeAttr("N", 2, concat_node);
// Shared init node
if (include_shared_init) {
if (!test_variable) {
CreateNode(shared_init_name, "NoOp", {}, &graph_def);
} else {
CreateNode(shared_init_name, "NoOp", {assign_node1, assign_node2},
&graph_def, true);
}
}
// Run the op.
GraphDef result;
TransformFuncContext context;
context.input_names = {"ids"};
context.output_names = {"gather1", "gather2"};
if (test_variable) {
context.params["input_checkpoint"] = {checkpoint_path};
}
if (shared_init_name != "group_deps") {
context.params["group_init_node"] = {shared_init_name};
}
TF_ASSERT_OK(SparsifyGather(graph_def, context, &result));
// Validation begins.
std::map<std::string, const NodeDef*> node_lookup;
MapNamesToNodes(result, &node_lookup);
// Check nodes.
EXPECT_EQ(0,
node_lookup.count("w1/part_1/Initializer/zeros/shape_as_tensor"));
EXPECT_EQ(0, node_lookup.count("w1/part_1/Initializer/zeros/Const"));
EXPECT_EQ(0, node_lookup.count("w1/part_1/Initializer/zeros"));
EXPECT_EQ(0, node_lookup.count("w1/part_1/Assign"));
EXPECT_EQ(0,
node_lookup.count("w2/part_1/Initializer/zeros/shape_as_tensor"));
EXPECT_EQ(0, node_lookup.count("w2/part_1/Initializer/zeros/Const"));
EXPECT_EQ(0, node_lookup.count("w2/part_1/Initializer/zeros"));
EXPECT_EQ(0, node_lookup.count("w2/part_1/Assign"));
EXPECT_EQ(1, node_lookup.count("ids"));
EXPECT_EQ("Const", node_lookup.at("ids")->op());
EXPECT_EQ(1, node_lookup.count(shared_init_name));
EXPECT_EQ("NoOp", node_lookup.at(shared_init_name)->op());
EXPECT_EQ(1, node_lookup.count("w1/part_1/indices"));
EXPECT_EQ("Const", node_lookup.at("w1/part_1/indices")->op());
Tensor expected_indices_tensor1(DT_INT64, TensorShape({3}));
test::FillValues<int64_t>(&expected_indices_tensor1, {0, 2, 3});
test::ExpectTensorEqual<int64_t>(
expected_indices_tensor1,
GetNodeTensorAttr(*(node_lookup.at("w1/part_1/indices")), "value"));
EXPECT_EQ(1, node_lookup.count("w1/part_1/values"));
EXPECT_EQ("Const", node_lookup.at("w1/part_1/values")->op());
Tensor expected_values_tensor1(DT_FLOAT, TensorShape({3}));
test::FillValues<float>(&expected_values_tensor1, {0.2, 1.2, 0.001});
test::ExpectTensorNear<float>(
expected_values_tensor1,
GetNodeTensorAttr(*(node_lookup.at("w1/part_1/values")), "value"),
1e-5);
EXPECT_EQ(1, node_lookup.count("w1/part_1/HashTable"));
EXPECT_EQ("HashTable", node_lookup.at("w1/part_1/HashTable")->op());
EXPECT_EQ(1, node_lookup.count("w1/part_1/InitializeTable"));
EXPECT_EQ("InitializeTable",
node_lookup.at("w1/part_1/InitializeTable")->op());
// Nodes in "gather1" scope.
EXPECT_EQ(1, node_lookup.count("gather1/LookupTableFind"));
EXPECT_EQ("LookupTableFind",
node_lookup.at("gather1/LookupTableFind")->op());
EXPECT_EQ(1, node_lookup.count("gather1/Const"));
EXPECT_EQ("Const", node_lookup.at("gather1/Const")->op());
Tensor expected_gather_default_tensor1(DT_FLOAT, TensorShape({}));
test::FillValues<float>(&expected_gather_default_tensor1, {0.0});
test::ExpectTensorNear<float>(
expected_gather_default_tensor1,
GetNodeTensorAttr(*(node_lookup.at("gather1/Const")), "value"), 1e-5);
EXPECT_EQ(1, node_lookup.count("gather1/ExpandDims/Const"));
EXPECT_EQ("Const", node_lookup.at("gather1/ExpandDims/Const")->op());
Tensor expected_expand_dims_tensor1(DT_INT32, TensorShape({}));
test::FillValues<int32_t>(&expected_expand_dims_tensor1, {-1});
test::ExpectTensorEqual<int32_t>(
expected_expand_dims_tensor1,
GetNodeTensorAttr(*(node_lookup.at("gather1/ExpandDims/Const")),
"value"));
EXPECT_EQ(1, node_lookup.count("gather1"));
EXPECT_EQ("ExpandDims", node_lookup.at("gather1")->op());
EXPECT_EQ(1, node_lookup.count("w2/part_1/indices"));
EXPECT_EQ("Const", node_lookup.at("w2/part_1/indices")->op());
Tensor expected_indices_tensor2(DT_INT64, TensorShape({3}));
test::FillValues<int64_t>(&expected_indices_tensor2, {0, 2, 3});
test::ExpectTensorEqual<int64_t>(
expected_indices_tensor2,
GetNodeTensorAttr(*(node_lookup.at("w2/part_1/indices")), "value"));
EXPECT_EQ(1, node_lookup.count("w2/part_1/values"));
EXPECT_EQ("Const", node_lookup.at("w2/part_1/values")->op());
Tensor expected_values_tensor2(DT_FLOAT, TensorShape({3}));
test::FillValues<float>(&expected_values_tensor2, {0.2, 1.2, 0.001});
test::ExpectTensorNear<float>(
expected_values_tensor2,
GetNodeTensorAttr(*(node_lookup.at("w2/part_1/values")), "value"),
1e-5);
EXPECT_EQ(1, node_lookup.count("w2/part_1/HashTable"));
EXPECT_EQ("HashTable", node_lookup.at("w2/part_1/HashTable")->op());
EXPECT_EQ(1, node_lookup.count("w2/part_1/InitializeTable"));
EXPECT_EQ("InitializeTable",
node_lookup.at("w2/part_1/InitializeTable")->op());
// Nodes in "gather2" scope.
EXPECT_EQ(1, node_lookup.count("gather2/LookupTableFind"));
EXPECT_EQ("LookupTableFind",
node_lookup.at("gather2/LookupTableFind")->op());
EXPECT_EQ(1, node_lookup.count("gather2/Const"));
EXPECT_EQ("Const", node_lookup.at("gather2/Const")->op());
Tensor expected_gather_default_tensor2(DT_FLOAT, TensorShape({}));
test::FillValues<float>(&expected_gather_default_tensor2, {0.0});
test::ExpectTensorNear<float>(
expected_gather_default_tensor2,
GetNodeTensorAttr(*(node_lookup.at("gather2/Const")), "value"), 1e-5);
EXPECT_EQ(1, node_lookup.count("gather2/ExpandDims/Const"));
EXPECT_EQ("Const", node_lookup.at("gather2/ExpandDims/Const")->op());
Tensor expected_expand_dims_tensor2(DT_INT32, TensorShape({}));
test::FillValues<int32_t>(&expected_expand_dims_tensor2, {-1});
test::ExpectTensorEqual<int32_t>(
expected_expand_dims_tensor2,
GetNodeTensorAttr(*(node_lookup.at("gather2/ExpandDims/Const")),
"value"));
EXPECT_EQ(1, node_lookup.count("gather2"));
EXPECT_EQ("ExpandDims", node_lookup.at("gather2")->op());
// Check connections
EXPECT_EQ("w1/part_1/HashTable",
node_lookup.at("w1/part_1/InitializeTable")->input(0));
EXPECT_EQ("w1/part_1/indices",
node_lookup.at("w1/part_1/InitializeTable")->input(1));
EXPECT_EQ("w1/part_1/values",
node_lookup.at("w1/part_1/InitializeTable")->input(2));
EXPECT_EQ("w2/part_1/HashTable",
node_lookup.at("w2/part_1/InitializeTable")->input(0));
EXPECT_EQ("w2/part_1/indices",
node_lookup.at("w2/part_1/InitializeTable")->input(1));
EXPECT_EQ("w2/part_1/values",
node_lookup.at("w2/part_1/InitializeTable")->input(2));
EXPECT_EQ("w1/part_1/HashTable",
node_lookup.at("gather1/LookupTableFind")->input(0));
EXPECT_EQ("ids", node_lookup.at("gather1/LookupTableFind")->input(1));
EXPECT_EQ("gather1/Const",
node_lookup.at("gather1/LookupTableFind")->input(2));
EXPECT_EQ("gather1/LookupTableFind", node_lookup.at("gather1")->input(0));
EXPECT_EQ("w2/part_1/HashTable",
node_lookup.at("gather2/LookupTableFind")->input(0));
EXPECT_EQ("ids", node_lookup.at("gather2/LookupTableFind")->input(1));
EXPECT_EQ("gather2/Const",
node_lookup.at("gather2/LookupTableFind")->input(2));
EXPECT_EQ("gather2/LookupTableFind", node_lookup.at("gather2")->input(0));
EXPECT_EQ(0, node_lookup.count("linear/concat/axis"));
EXPECT_EQ(0, node_lookup.count("concat/node"));
// Check control deps.
EXPECT_EQ(2, node_lookup.at(shared_init_name)->input_size());
EXPECT_NE(std::find(node_lookup.at(shared_init_name)->input().begin(),
node_lookup.at(shared_init_name)->input().end(),
"^w1/part_1/InitializeTable"),
node_lookup.at(shared_init_name)->input().end());
EXPECT_NE(std::find(node_lookup.at(shared_init_name)->input().begin(),
node_lookup.at(shared_init_name)->input().end(),
"^w2/part_1/InitializeTable"),
node_lookup.at(shared_init_name)->input().end());
}
void TestReadTensorSlice() {
const auto checkpoint_path =
io::JoinPath(testing::TmpDir(), "checkpoint_slice");
Tensor weights(DT_FLOAT, TensorShape({2, 1}));
test::FillValues<float>(&weights, {0.2, 0.000001});
BundleWriter writer(Env::Default(), checkpoint_path);
TF_ASSERT_OK(writer.AddSlice("w", TensorShape({4, 1}),
TensorSlice::ParseOrDie("0,2:0,1"), weights));
TF_ASSERT_OK(writer.Finish());
std::unique_ptr<BundleReader> reader(
new BundleReader(Env::Default(), checkpoint_path));
Tensor results;
TF_ASSERT_OK(
ReadTensorFromCheckpoint("w/part_0", reader, "4 1 0,2:0,1", &results));
test::ExpectTensorEqual<float>(weights, results);
}
};
TEST_F(SparsifyGatherTest, TestSinglePartition) {
TestSinglePartition(false, false, false, false);
TestSinglePartition(false, true, false, false);
TestSinglePartition(true, false, false, false);
TestSinglePartition(true, true, false, false);
TestSinglePartition(false, false, true, false);
TestSinglePartition(false, true, true, false);
TestSinglePartition(true, false, true, false);
TestSinglePartition(true, true, true, false);
TestSinglePartition(false, true, false, false, "shared_inits");
TestSinglePartition(true, true, false, false, "shared_inits");
TestSinglePartition(false, true, true, false, "shared_inits");
TestSinglePartition(true, true, true, false, "shared_inits");
TestSinglePartition(false, false, false, true);
TestSinglePartition(false, true, false, true);
TestSinglePartition(true, false, false, true);
TestSinglePartition(true, true, false, true);
TestSinglePartition(false, false, true, true);
TestSinglePartition(false, true, true, true);
TestSinglePartition(true, false, true, true);
TestSinglePartition(true, true, true, true);
TestSinglePartition(false, true, false, true, "shared_inits");
TestSinglePartition(true, true, false, true, "shared_inits");
TestSinglePartition(false, true, true, true, "shared_inits");
TestSinglePartition(true, true, true, true, "shared_inits");
}
TEST_F(SparsifyGatherTest, TestMultiPartition) {
TestMultiPartition(false, false, false);
TestMultiPartition(false, true, false);
TestMultiPartition(true, false, false);
TestMultiPartition(true, true, false);
TestMultiPartition(false, false, true);
TestMultiPartition(false, true, true);
TestMultiPartition(true, false, true);
TestMultiPartition(true, true, true);
TestMultiPartition(false, true, false, "shared_inits");
TestMultiPartition(true, true, false, "shared_inits");
TestMultiPartition(false, true, true, "shared_inits");
TestMultiPartition(true, true, true, "shared_inits");
}
TEST_F(SparsifyGatherTest, TestTensorSlice) { TestReadTensorSlice(); }
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/strip_unused_nodes.cc
+198
View File
@@ -0,0 +1,198 @@
/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "tensorflow/core/common_runtime/constant_folding.h"
#include "tensorflow/core/common_runtime/graph_constructor.h"
#include "tensorflow/core/graph/node_builder.h"
#include "tensorflow/core/graph/subgraph.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/tools/graph_transforms/fold_constants_lib.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
namespace {
absl::Status TypeForPlaceholder(const TransformFuncContext& context,
const std::string& node_name,
DataType* result) {
// If we don't find anything else, return float.
*result = DT_FLOAT;
// Check to see if we have been given a default for all placeholders.
if (context.params.count("type")) {
if (context.params.at("type").size() != 1) {
return absl::InvalidArgumentError(
"You must pass no more than one default 'type' to "
"strip_unused_nodes");
}
const std::string& type_string = context.params.at("type")[0];
if (!DataTypeFromString(type_string, result)) {
return absl::InvalidArgumentError(absl::StrCat(
"Couldn't understand type argument '", type_string, "'"));
}
}
// See if there's a particular type specified for this placeholder.
if (context.params.count("name") || context.params.count("type_for_name")) {
if (!context.params.count("name") ||
!context.params.count("type_for_name") ||
(context.params.at("type_for_name").size() !=
context.params.at("name").size())) {
return absl::InvalidArgumentError(
"You must pass a 'type_for_name' arg for every 'name', e.g. "
"strip_unused_nodes(name=foo, type_for_name=float, name=bar, "
"type_for_name=quint8");
}
const int name_count = context.params.at("name").size();
for (int i = 0; i < name_count; ++i) {
if (context.params.at("name")[i] == node_name) {
const std::string& type_string = context.params.at("type_for_name")[i];
if (!DataTypeFromString(type_string, result)) {
return absl::InvalidArgumentError(absl::StrCat(
"Couldn't understand type argument '", type_string, "'"));
}
}
}
}
return absl::OkStatus();
}
absl::Status ShapeForPlaceholder(const TransformFuncContext& context,
const std::string& node_name,
TensorShape* result) {
// If we don't find anything else, return scalar.
*result = {};
// Check to see if we have been given a default for all placeholders.
if (context.params.count("shape")) {
if (context.params.at("shape").size() != 1) {
return absl::InvalidArgumentError(
"You must pass no more than one default 'shape' to "
"strip_unused_nodes");
}
const std::string& shape_string = context.params.at("shape")[0];
TF_RETURN_IF_ERROR(TensorShapeFromString(shape_string, result));
}
// See if there's a particular type specified for this placeholder.
if (context.params.count("name") || context.params.count("shape_for_name")) {
if (!context.params.count("name") ||
!context.params.count("shape_for_name") ||
(context.params.at("shape_for_name").size() !=
context.params.at("name").size())) {
return absl::InvalidArgumentError(
"You must pass a 'shape_for_name' arg for every 'name', e.g. "
"strip_unused_nodes(name=foo, shape_for_name=\"2,2,1\", name=bar, "
"shape_for_name=\"1\"");
}
const int name_count = context.params.at("name").size();
for (int i = 0; i < name_count; ++i) {
if (context.params.at("name")[i] == node_name) {
const std::string& shape_string =
context.params.at("shape_for_name")[i];
TF_RETURN_IF_ERROR(TensorShapeFromString(shape_string, result));
}
}
}
return absl::OkStatus();
}
} // namespace
// Delete any nodes that don't contribute to the inference result.
absl::Status StripUnusedNodes(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def) {
std::set<std::string> required_nodes;
std::set<std::string> input_nodes;
for (const std::string& input : context.input_names) {
required_nodes.insert(NodeNameFromInput(input));
input_nodes.insert(NodeNameFromInput(input));
}
for (const std::string& output : context.output_names) {
required_nodes.insert(output);
}
std::map<std::string, const NodeDef*> node_lookup;
MapNamesToNodes(input_graph_def, &node_lookup);
std::vector<std::string> current_inputs;
current_inputs.reserve(context.output_names.size());
for (const std::string& output_name : context.output_names) {
current_inputs.push_back(NodeNameFromInput(output_name));
}
while (!current_inputs.empty()) {
std::set<std::string> next_inputs;
for (const std::string& current_input : current_inputs) {
required_nodes.insert(current_input);
if (input_nodes.count(current_input)) {
continue;
}
if (!node_lookup.count(current_input)) {
return absl::InvalidArgumentError(
absl::StrCat("Input node ", current_input, " not found in graph"));
}
const NodeDef* current_node = node_lookup[current_input];
for (const std::string& input_name : current_node->input()) {
std::string input_node_name = NodeNameFromInput(input_name);
if (!required_nodes.count(input_node_name)) {
next_inputs.insert(input_node_name);
}
}
}
current_inputs =
std::vector<std::string>(next_inputs.begin(), next_inputs.end());
}
GraphDef filtered_graph_def;
FilterGraphDef(input_graph_def,
[&](const NodeDef& node) {
return required_nodes.count(node.name()) > 0;
},
&filtered_graph_def);
output_graph_def->Clear();
for (const NodeDef& node : filtered_graph_def.node()) {
if (input_nodes.count(node.name())) {
NodeDef placeholder_node;
if (node.op() == "Placeholder") {
placeholder_node = node;
} else {
placeholder_node.set_op("Placeholder");
placeholder_node.set_name(node.name());
DataType type;
TF_RETURN_IF_ERROR(TypeForPlaceholder(context, node.name(), &type));
TensorShape shape;
TF_RETURN_IF_ERROR(ShapeForPlaceholder(context, node.name(), &shape));
SetNodeAttr("dtype", type, &placeholder_node);
SetNodeAttr("shape", shape, &placeholder_node);
}
*(output_graph_def->mutable_node()->Add()) = placeholder_node;
} else {
*(output_graph_def->mutable_node()->Add()) = node;
}
}
return absl::OkStatus();
}
REGISTER_GRAPH_TRANSFORM("strip_unused_nodes", StripUnusedNodes);
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/strip_unused_nodes_test.cc
+287
View File
@@ -0,0 +1,287 @@
/* Copyright 2015 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/cc/ops/const_op.h"
#include "tensorflow/cc/ops/image_ops.h"
#include "tensorflow/cc/ops/nn_ops.h"
#include "tensorflow/cc/ops/sendrecv_ops.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/core/framework/tensor_shape.pb.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
// Declare here, so we don't need a public header.
absl::Status StripUnusedNodes(const GraphDef& input_graph_def,
const TransformFuncContext& context,
GraphDef* output_graph_def);
class StripUnusedNodesTest : public ::testing::Test {
protected:
void TestSimpleAdd() {
GraphDef graph_def;
NodeDef* add_node = graph_def.add_node();
add_node->set_name("add_node");
add_node->set_op("Add");
add_node->add_input("a_node");
add_node->add_input("b_node");
NodeDef* a_node = graph_def.add_node();
a_node->set_name("a_node");
a_node->set_op("Const");
NodeDef* b_node = graph_def.add_node();
b_node->set_name("b_node");
b_node->set_op("Const");
NodeDef* c_node = graph_def.add_node();
c_node->set_name("c_node");
c_node->set_op("Const");
GraphDef result;
TF_ASSERT_OK(StripUnusedNodes(graph_def, {{}, {"add_node"}}, &result));
std::map<std::string, const NodeDef*> node_lookup;
MapNamesToNodes(result, &node_lookup);
EXPECT_EQ(1, node_lookup.count("add_node"));
EXPECT_EQ(1, node_lookup.count("a_node"));
EXPECT_EQ(1, node_lookup.count("b_node"));
EXPECT_EQ(0, node_lookup.count("c_node"));
}
void TestCommonAncestor() {
GraphDef graph_def;
NodeDef* add_node1 = graph_def.add_node();
add_node1->set_name("add_node1");
add_node1->set_op("Add");
add_node1->add_input("add_node2");
add_node1->add_input("add_node3");
NodeDef* add_node2 = graph_def.add_node();
add_node2->set_name("add_node2");
add_node2->set_op("Add");
add_node2->add_input("const_node1");
add_node2->add_input("const_node2");
NodeDef* add_node3 = graph_def.add_node();
add_node3->set_name("add_node3");
add_node3->set_op("Add");
add_node3->add_input("const_node1");
add_node3->add_input("const_node3");
NodeDef* const_node1 = graph_def.add_node();
const_node1->set_name("const_node1");
const_node1->set_op("Const");
NodeDef* const_node2 = graph_def.add_node();
const_node2->set_name("const_node2");
const_node2->set_op("Const");
NodeDef* const_node3 = graph_def.add_node();
const_node3->set_name("const_node3");
const_node3->set_op("Const");
NodeDef* dangling_input = graph_def.add_node();
dangling_input->set_name("dangling_input");
dangling_input->set_op("Const");
NodeDef* add_node4 = graph_def.add_node();
add_node4->set_name("add_node4");
add_node4->set_op("Add");
add_node4->add_input("add_node2");
add_node4->add_input("add_node3");
GraphDef result;
TF_ASSERT_OK(StripUnusedNodes(
graph_def, {{"dangling_input"}, {"add_node1"}}, &result));
std::map<std::string, const NodeDef*> node_lookup;
MapNamesToNodes(result, &node_lookup);
EXPECT_EQ(1, node_lookup.count("add_node1"));
EXPECT_EQ(1, node_lookup.count("add_node2"));
EXPECT_EQ(1, node_lookup.count("add_node3"));
EXPECT_EQ(0, node_lookup.count("add_node4"));
EXPECT_EQ(1, node_lookup.count("const_node1"));
EXPECT_EQ(1, node_lookup.count("const_node2"));
EXPECT_EQ(1, node_lookup.count("const_node3"));
EXPECT_EQ(0, node_lookup.count("const_node4"));
EXPECT_EQ(1, node_lookup.count("dangling_input"));
}
void TestSimplePlaceholder() {
GraphDef graph_def;
NodeDef* add_node = graph_def.add_node();
add_node->set_name("add_node");
add_node->set_op("Add");
add_node->add_input("mul_node");
add_node->add_input("a_node");
NodeDef* mul_node = graph_def.add_node();
mul_node->set_name("mul_node");
mul_node->set_op("Mul");
mul_node->add_input("b_node");
mul_node->add_input("c_node");
NodeDef* a_node = graph_def.add_node();
a_node->set_name("a_node");
a_node->set_op("Const");
NodeDef* b_node = graph_def.add_node();
b_node->set_name("b_node");
b_node->set_op("Const");
NodeDef* c_node = graph_def.add_node();
c_node->set_name("c_node");
c_node->set_op("Const");
GraphDef result;
TF_ASSERT_OK(
StripUnusedNodes(graph_def, {{"mul_node"}, {"add_node"}}, &result));
std::map<std::string, const NodeDef*> node_lookup;
MapNamesToNodes(result, &node_lookup);
EXPECT_EQ(1, node_lookup.count("add_node"));
EXPECT_EQ(1, node_lookup.count("mul_node"));
EXPECT_EQ("Placeholder", node_lookup["mul_node"]->op());
EXPECT_EQ(DT_FLOAT, node_lookup["mul_node"]->attr().at("dtype").type());
EXPECT_EQ(TensorShape({}),
TensorShape(node_lookup["mul_node"]->attr().at("shape").shape()));
EXPECT_EQ(1, node_lookup.count("a_node"));
EXPECT_EQ(0, node_lookup.count("b_node"));
EXPECT_EQ(0, node_lookup.count("c_node"));
}
void TestPlaceholderDefaultArgs() {
GraphDef graph_def;
NodeDef* add_node = graph_def.add_node();
add_node->set_name("add_node");
add_node->set_op("Add");
add_node->add_input("mul_node");
add_node->add_input("a_node");
NodeDef* mul_node = graph_def.add_node();
mul_node->set_name("mul_node");
mul_node->set_op("Mul");
mul_node->add_input("b_node");
mul_node->add_input("c_node");
NodeDef* a_node = graph_def.add_node();
a_node->set_name("a_node");
a_node->set_op("Const");
NodeDef* b_node = graph_def.add_node();
b_node->set_name("b_node");
b_node->set_op("Const");
NodeDef* c_node = graph_def.add_node();
c_node->set_name("c_node");
c_node->set_op("Const");
GraphDef result;
TF_ASSERT_OK(StripUnusedNodes(graph_def,
{{"mul_node"},
{"add_node"},
{{"type", {"int32"}}, {"shape", {"1,2,3"}}}},
&result));
std::map<std::string, const NodeDef*> node_lookup;
MapNamesToNodes(result, &node_lookup);
EXPECT_EQ(1, node_lookup.count("add_node"));
EXPECT_EQ(1, node_lookup.count("mul_node"));
EXPECT_EQ("Placeholder", node_lookup["mul_node"]->op());
EXPECT_EQ(DT_INT32, node_lookup["mul_node"]->attr().at("dtype").type());
EXPECT_EQ(TensorShape({1, 2, 3}),
TensorShape(node_lookup["mul_node"]->attr().at("shape").shape()));
EXPECT_EQ(1, node_lookup.count("a_node"));
EXPECT_EQ(0, node_lookup.count("b_node"));
EXPECT_EQ(0, node_lookup.count("c_node"));
}
void TestPlaceholderNamedArgs() {
GraphDef graph_def;
NodeDef* add_node = graph_def.add_node();
add_node->set_name("add_node");
add_node->set_op("Add");
add_node->add_input("mul_node");
add_node->add_input("a_node");
NodeDef* mul_node = graph_def.add_node();
mul_node->set_name("mul_node");
mul_node->set_op("Mul");
mul_node->add_input("b_node");
mul_node->add_input("c_node");
NodeDef* a_node = graph_def.add_node();
a_node->set_name("a_node");
a_node->set_op("Const");
NodeDef* b_node = graph_def.add_node();
b_node->set_name("b_node");
b_node->set_op("Const");
NodeDef* c_node = graph_def.add_node();
c_node->set_name("c_node");
c_node->set_op("Const");
GraphDef result;
TF_ASSERT_OK(StripUnusedNodes(graph_def,
{{"mul_node", "a_node"},
{"add_node"},
{{"name", {"a_node", "mul_node"}},
{"type_for_name", {"int64", "quint8"}},
{"shape_for_name", {"1,2", "1, 2, 3"}}}},
&result));
std::map<std::string, const NodeDef*> node_lookup;
MapNamesToNodes(result, &node_lookup);
EXPECT_EQ(1, node_lookup.count("add_node"));
EXPECT_EQ(1, node_lookup.count("mul_node"));
EXPECT_EQ("Placeholder", node_lookup["mul_node"]->op());
EXPECT_EQ(DT_QUINT8, node_lookup["mul_node"]->attr().at("dtype").type());
EXPECT_EQ(TensorShape({1, 2, 3}),
TensorShape(node_lookup["mul_node"]->attr().at("shape").shape()));
EXPECT_EQ(1, node_lookup.count("a_node"));
EXPECT_EQ("Placeholder", node_lookup["a_node"]->op());
EXPECT_EQ(DT_INT64, node_lookup["a_node"]->attr().at("dtype").type());
EXPECT_EQ(TensorShape({1, 2}),
TensorShape(node_lookup["a_node"]->attr().at("shape").shape()));
EXPECT_EQ(0, node_lookup.count("b_node"));
EXPECT_EQ(0, node_lookup.count("c_node"));
}
};
TEST_F(StripUnusedNodesTest, TestSimpleAdd) { TestSimpleAdd(); }
TEST_F(StripUnusedNodesTest, TestCommonAncestor) { TestCommonAncestor(); }
TEST_F(StripUnusedNodesTest, TestSimplePlaceholder) { TestSimplePlaceholder(); }
TEST_F(StripUnusedNodesTest, TestPlaceholderDefaultArgs) {
TestPlaceholderDefaultArgs();
}
TEST_F(StripUnusedNodesTest, TestPlaceholderNamedArgs) {
TestPlaceholderNamedArgs();
}
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/summarize_graph_main.cc
+344
View File
@@ -0,0 +1,344 @@
/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
// This program prints out a summary of a GraphDef file's contents, listing
// things that are useful for debugging and reusing the model it contains. For
// example it looks at the graph structure and op types to figure out likely
// input and output nodes, and shows which ops are used by the graph. To use it,
// run something like this:
//
// bazel build tensorflow/tools/graph_transforms:summarize_graph
// bazel-bin/tensorflow/tools/graph_transforms/summarize_graph \
// --in_graph=my_graph.pb
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.pb.h"
#include "tensorflow/core/lib/strings/str_util.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/util/command_line_flags.h"
#include "tensorflow/tools/graph_transforms/file_utils.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
namespace {
void PrintNodeInfo(const NodeDef* node) {
std::string shape_description = "None";
if (node->attr().count("shape")) {
TensorShapeProto shape_proto = node->attr().at("shape").shape();
absl::Status shape_status = PartialTensorShape::IsValidShape(shape_proto);
if (shape_status.ok()) {
shape_description = PartialTensorShape(shape_proto).DebugString();
} else {
shape_description = shape_status.message();
}
}
DataType dtype = DT_INVALID;
if (node->attr().count("dtype")) {
dtype = node->attr().at("dtype").type();
}
std::cout << "(name=" << node->name();
std::cout << ", type=" << DataTypeString(dtype) << "(" << dtype << ")";
std::cout << ", shape=" << shape_description << ") ";
}
void PrintBenchmarkUsage(const std::vector<const NodeDef*>& placeholders,
const std::vector<const NodeDef*>& variables,
const std::vector<const NodeDef*> outputs,
const std::string& graph_path) {
std::vector<const NodeDef*> all_inputs(placeholders);
all_inputs.insert(all_inputs.end(), variables.begin(), variables.end());
std::vector<std::string> input_layers;
std::vector<std::string> input_layer_types;
std::vector<std::string> input_layer_shapes;
for (const NodeDef* node : all_inputs) {
input_layers.push_back(node->name());
DataType dtype = DT_INVALID;
if (node->attr().count("dtype")) {
dtype = node->attr().at("dtype").type();
}
input_layer_types.push_back(DataTypeString(dtype));
std::vector<int64_t> sizes;
PartialTensorShape shape;
if (node->attr().count("shape")) {
TensorShapeProto shape_proto = node->attr().at("shape").shape();
if (PartialTensorShape::IsValid(shape_proto)) {
shape = PartialTensorShape(shape_proto);
}
}
std::string sizes_string;
if (shape.dims() == -1) {
// Unknown shapes can have -1 for dims, so leave these blank.
sizes_string = "";
} else {
sizes.reserve(shape.dims());
for (int i = 0; i < shape.dims(); ++i) {
sizes.push_back(shape.dim_size(i));
}
sizes_string = absl::StrJoin(sizes, ",");
}
input_layer_shapes.push_back(sizes_string);
}
std::vector<std::string> output_layers;
output_layers.reserve(outputs.size());
for (const NodeDef* node : outputs) {
output_layers.push_back(node->name());
}
std::string input_layer_value = absl::StrJoin(input_layers, ",");
std::string input_layer_type_value = absl::StrJoin(input_layer_types, ",");
std::string input_layer_shape_value = absl::StrJoin(input_layer_shapes, ":");
std::string output_layer_value = absl::StrJoin(output_layers, ",");
std::cout << "To use with tensorflow/tools/benchmark:benchmark_model try "
"these arguments:"
<< std::endl;
std::cout << "bazel run tensorflow/tools/benchmark:benchmark_model --";
std::cout << " --graph=" << graph_path;
std::cout << " --show_flops";
std::cout << " --input_layer=" << input_layer_value;
std::cout << " --input_layer_type=" << input_layer_type_value;
std::cout << " --input_layer_shape=" << input_layer_shape_value;
std::cout << " --output_layer=" << output_layer_value;
std::cout << std::endl;
}
absl::Status PrintStructure(const GraphDef& graph) {
GraphDef sorted_graph;
TF_RETURN_IF_ERROR(SortByExecutionOrder(graph, &sorted_graph));
for (const NodeDef& node : sorted_graph.node()) {
std::cout << node.name() << " (" << node.op() << "): ["
<< absl::StrJoin(node.input(), ", ") << "]";
if (node.op() == "Const") {
Tensor tensor;
if (node.attr().count("value") &&
tensor.FromProto(node.attr().at("value").tensor())) {
std::cout << ", value=" << tensor.DebugString();
} else {
LOG(WARNING) << "Decoding Tensor failed for node" << node.name();
}
}
std::cout << std::endl;
}
return absl::OkStatus();
}
absl::Status SummarizeGraph(const GraphDef& graph,
const std::string& graph_path,
bool print_structure) {
std::vector<const NodeDef*> placeholders;
std::vector<const NodeDef*> variables;
for (const NodeDef& node : graph.node()) {
if (node.op() == "Placeholder") {
placeholders.push_back(&node);
}
if (node.op() == "Variable" || node.op() == "VariableV2") {
variables.push_back(&node);
}
}
if (placeholders.empty()) {
std::cout << "No inputs spotted." << std::endl;
} else {
std::cout << "Found " << placeholders.size() << " possible inputs: ";
for (const NodeDef* node : placeholders) {
PrintNodeInfo(node);
}
std::cout << std::endl;
}
if (variables.empty()) {
std::cout << "No variables spotted." << std::endl;
} else {
std::cout << "Found " << variables.size() << " variables: ";
for (const NodeDef* node : variables) {
PrintNodeInfo(node);
}
std::cout << std::endl;
}
std::map<std::string, std::vector<const NodeDef*>> output_map;
MapNodesToOutputs(graph, &output_map);
std::vector<const NodeDef*> outputs;
std::unordered_set<std::string> unlikely_output_types = {
"Const", "Assign", "NoOp", "Placeholder"};
for (const NodeDef& node : graph.node()) {
if ((output_map.count(node.name()) == 0) &&
(unlikely_output_types.count(node.op()) == 0)) {
outputs.push_back(&node);
}
}
if (outputs.empty()) {
std::cout << "No outputs spotted." << std::endl;
} else {
std::cout << "Found " << outputs.size() << " possible outputs: ";
for (const NodeDef* node : outputs) {
std::cout << "(name=" << node->name();
std::cout << ", op=" << node->op() << ") ";
}
std::cout << std::endl;
}
int64_t const_parameter_count = 0;
int64_t variable_parameter_count = 0;
int control_edge_count = 0;
std::map<std::string, int> device_counts;
for (const NodeDef& node : graph.node()) {
for (const std::string& input : node.input()) {
if (input.substr(0, 1) == "^") {
++control_edge_count;
}
}
if (!node.device().empty()) {
++device_counts[node.device()];
}
if ((node.op() == "Const") || (node.op() == "Variable") ||
(node.op() == "VariableV2")) {
Tensor tensor;
if (node.attr().count("value") &&
tensor.FromProto(node.attr().at("value").tensor())) {
const size_t num_elements = tensor.NumElements();
if (node.op() == "Const") {
const_parameter_count += num_elements;
} else {
variable_parameter_count += num_elements;
}
} else {
LOG(WARNING) << "Decoding Tensor failed for node" << node.name();
}
}
}
std::cout << "Found " << const_parameter_count << " ("
<< strings::HumanReadableNum(const_parameter_count)
<< ") const parameters, " << variable_parameter_count << " ("
<< strings::HumanReadableNum(variable_parameter_count)
<< ") variable parameters, and " << control_edge_count
<< " control_edges" << std::endl;
if (!device_counts.empty()) {
for (const auto& device_info : device_counts) {
std::cout << device_info.second << " nodes assigned to device '"
<< device_info.first << "'";
}
}
std::vector<std::pair<std::string, std::string>> invalid_inputs;
FindInvalidInputs(graph, &invalid_inputs);
if (!invalid_inputs.empty()) {
for (const std::pair<std::string, std::string>& invalid_input :
invalid_inputs) {
std::cout << "Invalid input " << invalid_input.second << " for node "
<< invalid_input.first << std::endl;
}
return absl::InternalError(
"Invalid graph with inputs referring to nonexistent nodes");
}
std::map<std::string, int> op_counts;
for (const NodeDef& node : graph.node()) {
++op_counts[node.op()];
}
for (const FunctionDef& function : graph.library().function()) {
for (const NodeDef& node : function.node_def()) {
++op_counts[node.op()];
}
}
std::vector<std::pair<std::string, int>> op_counts_vec(op_counts.begin(),
op_counts.end());
std::sort(op_counts_vec.begin(), op_counts_vec.end(),
[](std::pair<std::string, int> a, std::pair<std::string, int> b) {
return (a.second > b.second);
});
std::cout << "Op types used: ";
bool is_first = true;
for (const std::pair<std::string, int>& op_count : op_counts_vec) {
if (!is_first) {
std::cout << ", ";
} else {
is_first = false;
}
std::cout << op_count.second << " " << op_count.first;
}
std::cout << std::endl;
PrintBenchmarkUsage(placeholders, variables, outputs, graph_path);
if (print_structure) {
TF_RETURN_IF_ERROR(PrintStructure(graph));
}
return absl::OkStatus();
}
int ParseFlagsAndSummarizeGraph(int argc, char* argv[]) {
std::string in_graph = "";
bool print_structure = false;
std::vector<Flag> flag_list = {
Flag("in_graph", &in_graph, "input graph file name"),
Flag("print_structure", &print_structure,
"whether to print the network connections of the graph"),
};
std::string usage = Flags::Usage(argv[0], flag_list);
const bool parse_result = Flags::Parse(&argc, argv, flag_list);
// We need to call this to set up global state for TensorFlow.
port::InitMain(argv[0], &argc, &argv);
if (!parse_result) {
LOG(ERROR) << usage;
return -1;
}
if (argc > 1) {
LOG(ERROR) << "Unknown argument " << argv[1] << ".\n" << usage;
return -1;
}
if (in_graph.empty()) {
LOG(ERROR) << "in_graph graph can't be empty.\n" << usage;
return -1;
}
GraphDef graph_def;
absl::Status load_status = LoadTextOrBinaryGraphFile(in_graph, &graph_def);
if (!load_status.ok()) {
LOG(ERROR) << "Loading graph '" << in_graph << "' failed with "
<< load_status.message();
LOG(ERROR) << usage;
return -1;
}
absl::Status summarize_result =
SummarizeGraph(graph_def, in_graph, print_structure);
if (!summarize_result.ok()) {
LOG(ERROR) << summarize_result.message() << "\n" << usage;
return -1;
}
return 0;
}
} // namespace
} // namespace graph_transforms
} // namespace tensorflow
int main(int argc, char* argv[]) {
return tensorflow::graph_transforms::ParseFlagsAndSummarizeGraph(argc, argv);
}
tensorflow/tools/graph_transforms/transform_graph.cc
+349
View File
@@ -0,0 +1,349 @@
/* Copyright 2016 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/tools/graph_transforms/transform_graph.h"
#include "tensorflow/core/framework/function.pb.h"
#include "tensorflow/core/lib/strings/scanner.h"
#include "tensorflow/core/lib/strings/str_util.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/util/command_line_flags.h"
#include "tensorflow/tools/graph_transforms/file_utils.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
#if !defined(PLATFORM_WINDOWS)
#include <pwd.h>
#include <unistd.h>
#endif
namespace tensorflow {
namespace graph_transforms {
using tensorflow::strings::Scanner;
absl::Status ParseTransformParameters(const std::string& transforms_string,
TransformParameters* params_list) {
params_list->clear();
enum {
TRANSFORM_NAME,
TRANSFORM_PARAM_NAME,
TRANSFORM_PARAM_VALUE,
} state = TRANSFORM_NAME;
absl::string_view remaining(transforms_string);
absl::string_view match;
absl::string_view transform_name;
absl::string_view parameter_name;
absl::string_view parameter_value;
TransformFuncParameters func_parameters;
while (!remaining.empty()) {
if (state == TRANSFORM_NAME) {
// Reset the list of parameters.
func_parameters.clear();
// Eat up any leading spaces.
Scanner(remaining).AnySpace().GetResult(&remaining, &match);
if (remaining.empty()) {
// Nothing remains after consuming trailing spaces.
// Consumed all transform parameter string without errors.
return absl::OkStatus();
}
// See if we have a valid transform name.
const bool found_transform_name =
Scanner(remaining)
.Many(Scanner::LETTER_DIGIT_UNDERSCORE)
.GetResult(&remaining, &transform_name);
if (!found_transform_name) {
return absl::InvalidArgumentError(
absl::StrCat("Looking for transform name, but found ", remaining));
}
if (Scanner(remaining).OneLiteral("(").GetResult(&remaining, &match)) {
state = TRANSFORM_PARAM_NAME;
} else {
// Add a transform with no parameters.
params_list->push_back({std::string(transform_name), func_parameters});
transform_name = "";
state = TRANSFORM_NAME;
}
} else if (state == TRANSFORM_PARAM_NAME) {
if (Scanner(remaining).OneLiteral(")").GetResult(&remaining, &match)) {
params_list->push_back({std::string(transform_name), func_parameters});
transform_name = "";
state = TRANSFORM_NAME;
} else {
// Eat up any leading spaces or commas.
Scanner(remaining).ZeroOrOneLiteral(",").GetResult(&remaining, &match);
Scanner(remaining).AnySpace().GetResult(&remaining, &match);
// See if we have a valid parameter name.
const bool found_parameter_name =
Scanner(remaining)
.Many(Scanner::LETTER_DIGIT_UNDERSCORE)
.GetResult(&remaining, &parameter_name);
if (!found_parameter_name) {
return absl::InvalidArgumentError(absl::StrCat(
"Looking for parameter name, but found ", remaining));
}
if (Scanner(remaining).OneLiteral("=").GetResult(&remaining, &match)) {
state = TRANSFORM_PARAM_VALUE;
} else {
return absl::InvalidArgumentError(
absl::StrCat("Looking for =, but found ", remaining));
}
}
} else if (state == TRANSFORM_PARAM_VALUE) {
bool found_parameter_value;
// Deal with quoted values.
if (Scanner(remaining).OneLiteral("\"").GetResult(&remaining, &match)) {
found_parameter_value =
Scanner(remaining).ScanEscapedUntil('"').GetResult(
&remaining, &parameter_value);
if (found_parameter_value) {
Scanner(remaining).OneLiteral("\"").GetResult(&remaining, &match);
}
} else {
// See if we have a valid parameter name.
found_parameter_value =
Scanner(remaining)
.Many(Scanner::LETTER_DIGIT_DASH_DOT_SLASH_UNDERSCORE)
.GetResult(&remaining, &parameter_value);
}
if (!found_parameter_value) {
return absl::InvalidArgumentError(
absl::StrCat("Looking for parameter name, but found ", remaining));
}
func_parameters[std::string(parameter_name)].emplace_back(
parameter_value);
// Eat up any trailing quotes.
Scanner(remaining).ZeroOrOneLiteral("\"").GetResult(&remaining, &match);
Scanner(remaining).ZeroOrOneLiteral("'").GetResult(&remaining, &match);
state = TRANSFORM_PARAM_NAME;
}
}
return absl::OkStatus();
}
std::string ExpandPath(const std::string& path_string) {
#if defined(PLATFORM_WINDOWS)
return path_string;
#else
if (path_string.empty() || path_string[0] != '~') {
return path_string;
}
const char* home = nullptr;
std::string::size_type prefix = path_string.find_first_of('/');
if (path_string.length() == 1 || prefix == 1) {
// The value of $HOME, e.g., ~/foo
home = getenv("HOME");
if (!home) {
// If HOME is not available, get uid
struct passwd* pw = getpwuid(getuid());
if (pw) {
home = pw->pw_dir;
}
}
} else {
// The value of ~user, e.g., ~user/foo
std::string user(path_string, 1, (prefix == std::string::npos)
? std::string::npos
: prefix - 1);
struct passwd* pw = getpwnam(user.c_str());
if (pw) {
home = pw->pw_dir;
}
}
if (!home) {
return path_string;
}
std::string path(home);
if (prefix == std::string::npos) {
return path;
}
if (path.length() == 0 || path[path.length() - 1] != '/') {
path += '/';
}
path += path_string.substr(prefix + 1);
return path;
#endif
}
int ParseFlagsAndTransformGraph(int argc, char* argv[], bool init_main) {
std::string in_graph_string = "";
std::string out_graph_string = "";
std::string inputs_string = "";
std::string outputs_string = "";
std::string transforms_string = "";
bool output_as_text = false;
std::vector<Flag> flag_list = {
Flag("in_graph", &in_graph_string, "input graph file name"),
Flag("out_graph", &out_graph_string, "output graph file name"),
Flag("inputs", &inputs_string, "inputs"),
Flag("outputs", &outputs_string, "outputs"),
Flag("transforms", &transforms_string, "list of transforms"),
Flag("output_as_text", &output_as_text,
"whether to write the graph in text protobuf format"),
};
std::string usage = Flags::Usage(argv[0], flag_list);
usage += "\nTransforms are:\n";
TransformRegistry* transform_registry = GetTransformRegistry();
for (const auto& pair : *transform_registry) {
usage += pair.first + "\n";
}
const bool parse_result = Flags::Parse(&argc, argv, flag_list);
// We need to call this to set up global state for TensorFlow.
if (init_main) {
port::InitMain(argv[0], &argc, &argv);
}
if (!parse_result) {
LOG(ERROR) << usage;
return -1;
}
if (argc > 1) {
LOG(ERROR) << "Unknown argument " << argv[1] << ".\n" << usage;
return -1;
}
if (in_graph_string.empty()) {
LOG(ERROR) << "in_graph graph can't be empty.\n" << usage;
return -1;
}
if (out_graph_string.empty()) {
LOG(ERROR) << "out_graph graph can't be empty.\n" << usage;
return -1;
}
if (transforms_string.empty()) {
LOG(ERROR) << "You must specify at least one transform.\n" << usage;
return -1;
}
std::string in_graph = ExpandPath(in_graph_string);
std::string out_graph = ExpandPath(out_graph_string);
std::vector<std::string> inputs = str_util::Split(inputs_string, ',');
std::vector<std::string> outputs = str_util::Split(outputs_string, ',');
TransformParameters transform_params;
absl::Status parse_status =
ParseTransformParameters(transforms_string, &transform_params);
if (!parse_status.ok()) {
LOG(ERROR) << "Failed to parse --transform argument, error was "
<< parse_status.message();
return -1;
}
if (transform_params.empty()) {
LOG(ERROR) << "You must specify at least one transform.\n" << usage;
return -1;
}
GraphDef graph_def;
absl::Status load_status = LoadTextOrBinaryGraphFile(in_graph, &graph_def);
if (!load_status.ok()) {
LOG(ERROR) << "Loading graph '" << in_graph_string << "' failed with "
<< load_status.message();
LOG(ERROR) << usage;
return -1;
}
absl::Status transform_result =
TransformGraph(inputs, outputs, transform_params, &graph_def);
if (!transform_result.ok()) {
LOG(ERROR) << transform_result.message();
LOG(ERROR) << usage;
return -1;
}
absl::Status save_status;
if (output_as_text) {
save_status = WriteTextProto(Env::Default(), out_graph, graph_def);
} else {
save_status = WriteBinaryProto(Env::Default(), out_graph, graph_def);
}
if (!save_status.ok()) {
LOG(ERROR) << "Saving graph '" << out_graph_string << "' failed with "
<< save_status.message();
return -1;
}
return 0;
}
absl::Status ShouldIgnoreErrors(const TransformFuncParameters& transform_params,
bool* ignore_errors) {
*ignore_errors = false;
if (transform_params.count("ignore_errors") &&
(!transform_params.at("ignore_errors").empty())) {
const std::string& ignore_errors_string =
absl::AsciiStrToLower(transform_params.at("ignore_errors").at(0));
if (ignore_errors_string == "true") {
*ignore_errors = true;
} else if (ignore_errors_string == "false") {
*ignore_errors = false;
} else {
return absl::InvalidArgumentError(
absl::StrCat("ignore_errors should be true or false, found ",
ignore_errors_string));
}
}
return absl::OkStatus();
}
absl::Status TransformGraph(const std::vector<std::string>& inputs,
const std::vector<std::string>& outputs,
const TransformParameters& transform_params,
GraphDef* graph_def) {
TransformRegistry* transform_registry = GetTransformRegistry();
for (const auto& transform_info : transform_params) {
const std::string& transform_name = transform_info.first;
if (transform_name.empty()) {
continue;
}
if (!transform_registry->count(transform_name)) {
return absl::InvalidArgumentError(
absl::StrCat("Transform '", transform_name, "' not recognized."));
}
LOG(INFO) << "Applying " << transform_name;
const TransformFunc& transform_func =
transform_registry->at(transform_name);
TransformFuncContext context;
context.input_names = inputs;
context.output_names = outputs;
context.params = transform_info.second;
bool ignore_errors;
TF_RETURN_IF_ERROR(
ShouldIgnoreErrors(transform_info.second, &ignore_errors));
GraphDef transformed_graph_def;
absl::Status transform_result =
transform_func(*graph_def, context, &transformed_graph_def);
if (!transform_result.ok()) {
if (ignore_errors) {
LOG(ERROR) << transform_name << ": Ignoring error "
<< transform_result.message();
transformed_graph_def = *graph_def;
} else {
return transform_result;
}
}
// Copy over the library from the original input graph.
*transformed_graph_def.mutable_library() = graph_def->library();
TF_RETURN_IF_ERROR(IsGraphValid(transformed_graph_def));
*graph_def = transformed_graph_def;
}
return absl::OkStatus();
}
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/transform_graph.h
+50
View File
@@ -0,0 +1,50 @@
/* Copyright 2015 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_TOOLS_GRAPH_TRANSFORMS_TRANSFORM_GRAPH_H_
#define TENSORFLOW_TOOLS_GRAPH_TRANSFORMS_TRANSFORM_GRAPH_H_
#include <vector>
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
// Convenience function to handle argument parsing for the command line tool.
// If init_main is false, we're testing so don't call core initialization.
int ParseFlagsAndTransformGraph(int argc, char* argv[], bool init_main);
// Handles converting the transforms string into transform names and their
// arguments.
typedef std::vector<std::pair<std::string, TransformFuncParameters>>
TransformParameters;
absl::Status ParseTransformParameters(const std::string& transforms_string,
TransformParameters* params_list);
// Applies a series of transformations to the GraphDef. These transforms are
// defined by modules that call REGISTER_GRAPH_TRANSFORM() to associate a
// function with a name string.
absl::Status TransformGraph(const std::vector<std::string>& inputs,
const std::vector<std::string>& outputs,
const TransformParameters& transform_params,
GraphDef* graph_def);
} // namespace graph_transforms
} // namespace tensorflow
#endif // TENSORFLOW_TOOLS_GRAPH_TRANSFORMS_TRANSFORM_GRAPH_H_
tensorflow/tools/graph_transforms/transform_graph_main.cc
+53
View File
@@ -0,0 +1,53 @@
/* Copyright 2016 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.
==============================================================================*/
// Tool that applies a series of transformations to a frozen GraphDef file.
// It takes a flexible list of transforms either on the command line, and runs
// those on the incoming graph to produce the result. This allows you to build a
// processing pipeline when preparing models for deployment.
//
// bazel build tensorflow/tools/graph_transforms/fold_constants_tool &&
// bazel-bin/tensorflow/tools/graph_transforms/fold_constants_tool \
// --in_graph=graph_def.pb \
// --out_graph=transformed_graph_def.pb \
// --inputs=input1,input2 \
// --outputs=output1,output2 \
// --transforms="fold_constants order_nodes"
//
// Parameters:
// in_graph - name of a file with a frozen GraphDef proto in binary format.
// out_graph - name of the output file to save the transformed version to.
// inputs - layer names of the nodes that will be fed data.
// outputs - layer names of the nodes that will be read from after running.
// transforms - space-separated names of the transforms to apply.
//
// List of implemented transforms:
// fold_constants - Merges constant expression subgraphs into single constants,
// which can help reduce the number of ops and make subsequent transforms
// optimizations more effective.
// order_nodes - Sorts the GraphDef nodes in execution order, which can help
// simple inference engines that want to avoid complexity in their executors.
#include "tensorflow/core/lib/strings/str_util.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/init_main.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/util/command_line_flags.h"
#include "tensorflow/tools/graph_transforms/transform_graph.h"
int main(int argc, char* argv[]) {
return tensorflow::graph_transforms::ParseFlagsAndTransformGraph(argc, argv,
true);
}
tensorflow/tools/graph_transforms/transform_graph_test.cc
+255
View File
@@ -0,0 +1,255 @@
/* Copyright 2015 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/tools/graph_transforms/transform_graph.h"
#include "tensorflow/cc/ops/const_op.h"
#include "tensorflow/cc/ops/image_ops.h"
#include "tensorflow/cc/ops/nn_ops.h"
#include "tensorflow/cc/ops/sendrecv_ops.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/lib/io/path.h"
#include "tensorflow/core/lib/strings/str_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/tools/graph_transforms/transform_utils.h"
namespace tensorflow {
namespace graph_transforms {
// Declared here so we don't have to expose it in the public header.
absl::Status ShouldIgnoreErrors(const TransformFuncParameters& transform_params,
bool* ignore_errors);
namespace {
absl::Status test_empty_graph_transform(const GraphDef& graph_def,
const TransformFuncContext& context,
GraphDef* result) {
result->Clear();
return absl::OkStatus();
}
} // namespace
REGISTER_GRAPH_TRANSFORM("test_empty_graph_transform",
test_empty_graph_transform);
class TransformGraphTest : public ::testing::Test {
protected:
void TestConstantFolding() {
auto root = tensorflow::Scope::NewRootScope();
using namespace ::tensorflow::ops; // NOLINT(build/namespaces)
const int width = 100;
Tensor a_data(DT_FLOAT, TensorShape({width}));
test::FillIota<float>(&a_data, 1.0f);
Output a_const =
Const(root.WithOpName("a_expect_removed"), Input::Initializer(a_data));
Tensor b_data(DT_FLOAT, TensorShape({width}));
test::FillIota<float>(&b_data, 1.0f);
Output b_const =
Const(root.WithOpName("b_expect_removed"), Input::Initializer(b_data));
Output add = Add(root.WithOpName("add_expect_removed"), a_const, b_const);
Output placeholder =
Placeholder(root.WithOpName("placeholder_expect_remains"), DT_FLOAT);
Output mul =
Mul(root.WithOpName("output_expect_remains"), add, placeholder);
GraphDef graph_def;
TF_ASSERT_OK(root.ToGraphDef(&graph_def));
std::string graph_def_serialized;
graph_def.SerializeToString(&graph_def_serialized);
const std::string dir = testing::TmpDir();
const std::string in_filename_pb = io::JoinPath(dir, "in_graphdef.pb");
const std::string out_filename_pb = io::JoinPath(dir, "out_graphdef.pb");
TF_ASSERT_OK(WriteStringToFile(Env::Default(), in_filename_pb,
graph_def_serialized));
std::vector<std::string> args = {"some_binary",
"--in_graph=" + in_filename_pb,
"--out_graph=" + out_filename_pb,
"--inputs=placeholder_expect_remains",
"--outputs=output_expect_remains",
"--transforms=fold_constants"};
const int argc = 6;
EXPECT_EQ(argc, args.size());
char* argv[argc];
std::vector<char*> char_strings;
for (int i = 0; i < argc; ++i) {
std::string arg = args[i];
char* char_string = new char[arg.size() + 1];
std::copy_n(arg.c_str(), arg.size() + 1, char_string);
argv[i] = char_string;
char_strings.push_back(char_string);
}
ParseFlagsAndTransformGraph(argc, argv, false);
for (char* char_string : char_strings) {
delete[] char_string;
}
GraphDef out_graph_def;
TF_EXPECT_OK(
ReadBinaryProto(Env::Default(), out_filename_pb, &out_graph_def));
std::map<std::string, const NodeDef*> out_node_map;
graph_transforms::MapNamesToNodes(out_graph_def, &out_node_map);
for (const NodeDef& node : out_graph_def.node()) {
const int occurrence_count = out_node_map.count(node.name());
if (absl::EndsWith(node.name(), "expect_removed")) {
EXPECT_EQ(0, occurrence_count) << "node.name()=" << node.name();
}
if (absl::EndsWith(node.name(), "expect_remains")) {
EXPECT_EQ(1, occurrence_count) << "node.name()=" << node.name();
}
}
}
void TestTransformRegistration() {
auto root = tensorflow::Scope::NewRootScope();
using namespace ::tensorflow::ops; // NOLINT(build/namespaces)
Output placeholder =
Placeholder(root.WithOpName("placeholder_expect_remains"), DT_FLOAT);
GraphDef graph_def;
TF_ASSERT_OK(root.ToGraphDef(&graph_def));
EXPECT_EQ(1, graph_def.node().size());
TF_ASSERT_OK(TransformGraph({}, {}, {{"test_empty_graph_transform", {}}},
&graph_def));
EXPECT_EQ(0, graph_def.node().size());
TF_ASSERT_OK(root.ToGraphDef(&graph_def));
absl::Status no_such_status =
TransformGraph({}, {}, {{"test_no_such_transform", {}}}, &graph_def);
EXPECT_TRUE(absl::StrContains(no_such_status.ToString(), "not recognized"));
}
void TestParseTransformParameters() {
TransformParameters params_list;
TF_EXPECT_OK(ParseTransformParameters("foo", &params_list));
EXPECT_EQ(1, params_list.size());
EXPECT_EQ("foo", params_list[0].first);
EXPECT_TRUE(params_list[0].second.empty());
TF_EXPECT_OK(ParseTransformParameters("foo bar", &params_list));
EXPECT_EQ(2, params_list.size());
EXPECT_EQ("foo", params_list[0].first);
EXPECT_TRUE(params_list[0].second.empty());
EXPECT_EQ("bar", params_list[1].first);
EXPECT_TRUE(params_list[1].second.empty());
TF_EXPECT_OK(ParseTransformParameters("foo() bar()", &params_list));
EXPECT_EQ(2, params_list.size());
EXPECT_EQ("foo", params_list[0].first);
EXPECT_TRUE(params_list[0].second.empty());
EXPECT_EQ("bar", params_list[1].first);
EXPECT_TRUE(params_list[1].second.empty());
TF_EXPECT_OK(
ParseTransformParameters("foo(bob_something=sue)", &params_list));
EXPECT_EQ(1, params_list.size());
EXPECT_EQ("foo", params_list[0].first);
EXPECT_EQ(1, params_list[0].second.count("bob_something"));
EXPECT_EQ(1, params_list[0].second["bob_something"].size());
EXPECT_EQ("sue", params_list[0].second["bob_something"][0]);
TF_EXPECT_OK(ParseTransformParameters("bar(a=1, b=2, a=3)", &params_list));
EXPECT_EQ(1, params_list.size());
EXPECT_EQ("bar", params_list[0].first);
EXPECT_EQ(1, params_list[0].second.count("a"));
EXPECT_EQ(2, params_list[0].second["a"].size());
EXPECT_EQ("1", params_list[0].second["a"][0]);
EXPECT_EQ("3", params_list[0].second["a"][1]);
EXPECT_EQ(1, params_list[0].second.count("b"));
EXPECT_EQ(1, params_list[0].second["b"].size());
EXPECT_EQ("2", params_list[0].second["b"][0]);
TF_EXPECT_OK(ParseTransformParameters("bar(a=\"1\", b=\"1,2,3\", a=3)",
&params_list));
EXPECT_EQ(1, params_list.size());
EXPECT_EQ("bar", params_list[0].first);
EXPECT_EQ(1, params_list[0].second.count("a"));
EXPECT_EQ(2, params_list[0].second["a"].size());
EXPECT_EQ("1", params_list[0].second["a"][0]);
EXPECT_EQ("3", params_list[0].second["a"][1]);
EXPECT_EQ(1, params_list[0].second.count("b"));
EXPECT_EQ(1, params_list[0].second["b"].size());
EXPECT_EQ("1,2,3", params_list[0].second["b"][0]);
}
void TestParseEscapedNewline() {
// This sequence of characters caused an infinite loop in the parser, which
// is responsible for the hang mentioned in
// https://github.com/tensorflow/tensorflow/issues/7150
TransformParameters params_list;
ParseTransformParameters("\\\n", &params_list).IgnoreError();
EXPECT_EQ(0, params_list.size());
}
void TestParseExtraSpaces() {
TransformParameters params_list;
ParseTransformParameters(" ", &params_list).IgnoreError();
EXPECT_EQ(0, params_list.size());
TF_EXPECT_OK(ParseTransformParameters(" foo bar \\\n", &params_list));
EXPECT_EQ(2, params_list.size());
EXPECT_EQ("foo", params_list[0].first);
EXPECT_TRUE(params_list[0].second.empty());
EXPECT_EQ("bar", params_list[1].first);
EXPECT_TRUE(params_list[1].second.empty());
}
void TestShouldIgnoreErrors() {
bool ignore_errors;
TF_EXPECT_OK(
ShouldIgnoreErrors({{"ignore_errors", {"true"}}}, &ignore_errors));
EXPECT_TRUE(ignore_errors);
TF_EXPECT_OK(
ShouldIgnoreErrors({{"ignore_errors", {"false"}}}, &ignore_errors));
EXPECT_FALSE(ignore_errors);
TF_EXPECT_OK(ShouldIgnoreErrors({}, &ignore_errors));
EXPECT_FALSE(ignore_errors);
EXPECT_FALSE(
ShouldIgnoreErrors({{"ignore_errors", {"foo"}}}, &ignore_errors).ok());
}
};
TEST_F(TransformGraphTest, TestConstantFolding) { TestConstantFolding(); }
TEST_F(TransformGraphTest, TestTransformRegistration) {
TestTransformRegistration();
}
TEST_F(TransformGraphTest, TestParseTransformParameters) {
TestParseTransformParameters();
}
TEST_F(TransformGraphTest, TestParseEscapedNewline) {
TestParseEscapedNewline();
}
TEST_F(TransformGraphTest, TestShouldIgnoreErrors) { TestShouldIgnoreErrors(); }
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/transform_utils.cc
+725
View File
@@ -0,0 +1,725 @@
/* Copyright 2015 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/tools/graph_transforms/transform_utils.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/lib/hash/hash.h"
#include "tensorflow/core/lib/strings/numbers.h"
#include "tensorflow/core/lib/strings/str_util.h"
namespace tensorflow {
namespace graph_transforms {
namespace {
inline bool IsMerge(const NodeDef& node_def) {
return node_def.op() == "Merge" || node_def.op() == "RefMerge" ||
node_def.op() == "_XlaMerge";
}
void RecordMatchedNodes(const NodeMatch& match,
std::set<std::string>* matched_nodes) {
matched_nodes->insert(match.node.name());
for (const NodeMatch& input_match : match.inputs) {
RecordMatchedNodes(input_match, matched_nodes);
}
}
inline uint64_t Hash64String(const std::string& input) {
return Hash64(input.data(), input.size());
}
} // namespace
void MatchedNodesAsArray(const NodeMatch& match, std::vector<NodeDef>* result) {
std::set<std::string> found_nodes;
std::vector<NodeMatch> current_matches = {match};
while (!current_matches.empty()) {
std::vector<NodeMatch> next_matches;
for (const NodeMatch& current_match : current_matches) {
if (found_nodes.count(current_match.node.name())) {
continue;
}
found_nodes.insert(current_match.node.name());
result->push_back(current_match.node);
for (const NodeMatch& input_match : current_match.inputs) {
next_matches.push_back(input_match);
}
}
current_matches = next_matches;
}
}
void MapNamesToNodes(const GraphDef& graph_def,
std::map<std::string, const NodeDef*>* result) {
for (const NodeDef& node : graph_def.node()) {
(*result)[node.name()] = &node;
}
}
void MapNodesToOutputs(
const GraphDef& graph_def,
std::map<std::string, std::vector<const NodeDef*>>* result) {
std::map<std::string, const NodeDef*> node_map;
MapNamesToNodes(graph_def, &node_map);
for (const NodeDef& node : graph_def.node()) {
for (const std::string& input : node.input()) {
std::string input_node_name = NodeNameFromInput(input);
(*result)[input_node_name].push_back(&node);
}
}
}
void NodeNamePartsFromInput(const std::string& input_name, std::string* prefix,
std::string* node_name, std::string* suffix) {
std::vector<std::string> input_parts = str_util::Split(input_name, ':');
if (input_parts.size() < 2) {
*suffix = "";
} else {
*suffix = ":" + input_parts[1];
}
absl::string_view node_name_piece(input_parts[0]);
if (absl::ConsumePrefix(&node_name_piece, "^")) {
*prefix = "^";
} else {
*prefix = "";
}
*node_name = std::string(node_name_piece);
}
std::string NodeNameFromInput(const std::string& input_name) {
std::string prefix;
std::string node_name;
std::string suffix;
NodeNamePartsFromInput(input_name, &prefix, &node_name, &suffix);
return node_name;
}
std::string CanonicalInputName(const std::string& input_name) {
std::string prefix;
std::string node_name;
std::string suffix;
NodeNamePartsFromInput(input_name, &prefix, &node_name, &suffix);
if (suffix.empty()) {
suffix = ":0";
}
return prefix + node_name + suffix;
}
uint64_t HashNodeDef(const NodeDef& node) {
uint64_t hash = Hash64String(node.op());
hash = Hash64Combine(hash, Hash64String(node.name()));
for (const std::string& input : node.input()) {
hash = Hash64Combine(hash, Hash64String(CanonicalInputName(input)));
}
hash = Hash64Combine(hash, Hash64String(node.device()));
std::vector<std::string> attr_names;
attr_names.reserve(node.attr().size());
for (const auto& attr : node.attr()) {
attr_names.push_back(attr.first);
}
std::sort(attr_names.begin(), attr_names.end());
std::string attr_serialized;
for (const std::string& attr_name : attr_names) {
auto attr = node.attr().at(attr_name);
attr.SerializeToString(&attr_serialized);
hash = Hash64Combine(hash, Hash64String(attr_serialized));
}
return hash;
}
void AddNodeInput(const std::string& input_name, NodeDef* node) {
*(node->mutable_input()->Add()) = input_name;
}
void CopyNodeAttr(const NodeDef& source, const std::string& source_key,
const std::string& dest_key, NodeDef* dest) {
CHECK_NE(0, source.attr().count(source_key))
<< "No key '" << source_key << "' found in " << source.DebugString();
(*(dest->mutable_attr()))[dest_key] = source.attr().at(source_key);
}
Tensor GetNodeTensorAttr(const NodeDef& node, const std::string& key) {
TensorProto tensor_proto = node.attr().at(key).tensor();
Tensor tensor;
CHECK(tensor.FromProto(tensor_proto));
return tensor;
}
void FilterGraphDef(const GraphDef& input_graph_def,
std::function<bool(const NodeDef&)> selector,
GraphDef* output_graph_def) {
output_graph_def->mutable_node()->Clear();
for (const NodeDef& node : input_graph_def.node()) {
if (selector(node)) {
*output_graph_def->mutable_node()->Add() = node;
}
}
}
void RemoveAttributes(const GraphDef& input_graph_def,
const std::vector<std::string>& attributes,
GraphDef* output_graph_def) {
output_graph_def->mutable_node()->Clear();
for (const NodeDef& node : input_graph_def.node()) {
NodeDef* new_node = output_graph_def->mutable_node()->Add();
*new_node = node;
for (const std::string& attribute : attributes) {
new_node->mutable_attr()->erase(attribute);
}
}
}
absl::Status SortByExecutionOrder(const GraphDef& input_graph_def,
GraphDef* output_graph_def) {
const int num_nodes = input_graph_def.node_size();
std::vector<int> ready;
std::vector<int> pending_count;
pending_count.reserve(num_nodes);
std::vector<absl::InlinedVector<int, 4UL>> outputs(num_nodes);
std::map<std::string, int> name_index;
for (int i = 0; i < input_graph_def.node_size(); ++i) {
const NodeDef& node(input_graph_def.node(i));
name_index[node.name()] = i;
}
// Parse the inputs for each node.
for (int n = 0; n < num_nodes; ++n) {
const NodeDef& node_def(input_graph_def.node(n));
if (IsMerge(node_def)) {
// for merge only wait for one non-control input.
int32_t num_control_edges = 0;
for (int i = 0; i < node_def.input_size(); ++i) {
if (absl::StartsWith(node_def.input(i), "^")) {
num_control_edges++;
}
}
pending_count.push_back(num_control_edges + 1);
} else {
pending_count.push_back(node_def.input_size());
}
if (node_def.input_size() == 0) {
ready.push_back(n);
continue;
}
for (int i = 0; i < node_def.input_size(); ++i) {
const std::string& input_name = node_def.input(i);
const std::string& input_node_name = NodeNameFromInput(input_name);
if (!name_index.count(input_node_name)) {
return absl::InvalidArgumentError(
absl::StrCat("Node '", node_def.name(), "': Unknown input node '",
node_def.input(i), "'"));
}
outputs[name_index[input_node_name]].push_back(n);
}
}
int processed = 0;
output_graph_def->Clear();
// Process the NodeDefs in topological order.
// Code above sets this up by filling in ready_ with nodes that have no
// inputs, pending_counts_ with the number of inputs for each node and
// outputs_ with the outputs of each node.
while (!ready.empty()) {
int o = ready.back();
ready.pop_back();
++processed;
const NodeDef& node_def(input_graph_def.node(o));
*output_graph_def->mutable_node()->Add() = node_def;
// Update pending_count for outputs.
for (size_t i = 0; i < outputs[o].size(); ++i) {
const int output = outputs[o][i];
pending_count[output]--;
if (pending_count[output] == 0) {
ready.push_back(output);
}
}
}
if (processed < num_nodes) {
LOG(WARNING) << "IN " << __func__ << (num_nodes - processed)
<< " NODES IN A CYCLE";
for (int64_t i = 0; i < num_nodes; i++) {
if (pending_count[i] != 0) {
LOG(WARNING) << "PENDING: " << SummarizeNodeDef(input_graph_def.node(i))
<< "WITH PENDING COUNT = " << pending_count[i];
}
}
return absl::InvalidArgumentError(
absl::StrCat(num_nodes - processed, " nodes in a cycle"));
}
return absl::OkStatus();
}
std::string OpTypePattern::DebugString() const {
std::string result = "{" + op + ", {";
for (const OpTypePattern& input : inputs) {
result += input.DebugString() + ",";
}
result += "}}";
return result;
}
std::string NodeMatch::DebugString() const {
std::string result = "{";
result += node.DebugString();
result += ", {";
for (const NodeMatch& input : inputs) {
result += input.DebugString() + ",";
}
result += "}}";
return result;
}
GraphMatcher::GraphMatcher(const GraphDef& graph_def) {
SortByExecutionOrder(graph_def, &graph_def_).IgnoreError();
MapNamesToNodes(graph_def_, &node_map_);
}
absl::Status GraphMatcher::GetOpTypeMatches(const OpTypePattern& pattern,
std::vector<NodeMatch>* matches) {
std::set<std::string> matched_nodes;
for (const NodeDef& node : graph_def_.node()) {
// Skip any nodes that are already part of a match.
if (matched_nodes.count(node.name())) {
continue;
}
NodeMatch match;
if (DoesOpTypeMatch(node, pattern, matched_nodes, &match)) {
RecordMatchedNodes(match, &matched_nodes);
matches->push_back(match);
}
}
return absl::OkStatus();
}
bool GraphMatcher::DoesOpTypeMatch(
const NodeDef& node, const OpTypePattern& pattern,
const std::set<std::string>& previously_matched_nodes, NodeMatch* match) {
VLOG(1) << "Looking at node " << node.DebugString();
VLOG(1) << "pattern=" << pattern.DebugString();
VLOG(1) << "match=" << match->DebugString();
if (previously_matched_nodes.count(node.name())) {
VLOG(1) << "node " << node.name() << " has been previously matched";
return false;
}
bool pattern_matched = false;
if (pattern.op == "*") {
pattern_matched = true;
} else {
std::vector<std::string> pattern_ops = str_util::Split(pattern.op, '|');
for (const std::string& pattern_op : pattern_ops) {
if (node.op() == pattern_op) {
pattern_matched = true;
}
}
}
if (!pattern_matched) {
VLOG(1) << "node.op() != pattern.op()";
return false;
}
match->node = node;
// Ignore any control inputs for pattern-matching purposes
std::vector<std::string> non_control_inputs;
for (const std::string& input : node.input()) {
if (!input.empty() && (input[0] != '^')) {
non_control_inputs.push_back(input);
}
}
if (pattern.inputs.empty()) {
// If there are no inputs, assume that's the end of the pattern.
return true;
}
if (non_control_inputs.size() != pattern.inputs.size()) {
VLOG(1) << "non_control_inputs.size() != pattern.inputs.size()";
return false;
}
for (int i = 0; i < pattern.inputs.size(); ++i) {
const std::string& input_node_name =
NodeNameFromInput(non_control_inputs[i]);
const NodeDef& input_node = *(node_map_[input_node_name]);
const OpTypePattern& input_pattern = pattern.inputs[i];
match->inputs.push_back(NodeMatch());
NodeMatch* input_match = &(match->inputs.back());
if (!DoesOpTypeMatch(input_node, input_pattern, previously_matched_nodes,
input_match)) {
return false;
}
}
return true;
}
absl::Status ReplaceMatchingOpTypes(
const GraphDef& input_graph_def, const OpTypePattern& pattern,
const std::function<absl::Status(
const NodeMatch&, const std::set<std::string>&,
const std::set<std::string>&, std::vector<NodeDef>*)>& node_generator,
const ReplaceMatchingOpTypesOptions& options, GraphDef* output_graph_def) {
// Start off by retrieving all the matching subgraphs.
GraphMatcher matcher(input_graph_def);
std::vector<NodeMatch> matches;
TF_RETURN_IF_ERROR(matcher.GetOpTypeMatches(pattern, &matches));
// Do some housekeeping so we can easily look up the resulting matches given
// a node name.
std::set<std::string> matched_nodes;
std::map<std::string, const NodeMatch*> matches_by_head_name;
for (const NodeMatch& match : matches) {
matches_by_head_name[match.node.name()] = &match;
RecordMatchedNodes(match, &matched_nodes);
}
std::map<std::string, std::vector<const NodeDef*>> outputs_map;
MapNodesToOutputs(input_graph_def, &outputs_map);
// Go through all the nodes in the input graph, see if they are part of a
// match or if they can be left untouched.
output_graph_def->Clear();
for (const NodeDef& input_node : input_graph_def.node()) {
if (matches_by_head_name.count(input_node.name())) {
// This node is the beginning of a match, so call the replacement function
// after setting up some information it will need.
const NodeMatch* match = matches_by_head_name[input_node.name()];
std::vector<NodeDef> matched_nodes_array;
MatchedNodesAsArray(*match, &matched_nodes_array);
// This tells us whether a node is part of the current match.
std::set<std::string> matched_nodes_lookup;
for (const NodeDef& matched_node : matched_nodes_array) {
matched_nodes_lookup.insert(matched_node.name());
}
// These are helper arrays that the replacement function can use to tell
// whether it can safely remove an internal node (because nothing outside
// of the match uses it) or whether external nodes depend on it.
std::set<std::string> input_nodes;
std::set<std::string> output_nodes;
for (const NodeDef& matched_node : matched_nodes_array) {
// Look through all of this node's inputs, and if any of them come from
// outside the match, then this should be noted as one of the external
// inputs of the subgraph.
for (const std::string& input_name : matched_node.input()) {
std::string input_node_name = NodeNameFromInput(input_name);
if (!matched_nodes_lookup.count(input_node_name)) {
input_nodes.insert(matched_node.name());
}
}
// Do a reverse input lookup, to see which other nodes use the current
// one as an input. If any of those nodes are outside the match
// subgraph, then the current node is marked as an output node that
// shouldn't be removed.
if (outputs_map.count(matched_node.name())) {
for (const NodeDef* dependent_node :
outputs_map[matched_node.name()]) {
if (!matched_nodes_lookup.count(dependent_node->name())) {
output_nodes.insert(matched_node.name());
}
}
}
}
// Call the generator function and add all the returned nodes to the
// graph.
std::vector<NodeDef> new_nodes;
TF_RETURN_IF_ERROR(
node_generator(*match, input_nodes, output_nodes, &new_nodes));
std::set<std::string> new_node_names;
for (const NodeDef& new_node : new_nodes) {
new_node_names.insert(new_node.name());
}
// Check to make sure the generator function preserved all of the nodes
// that are used elsewhere in the graph, and add them back in if not.
bool abort_replacement = false;
if (!options.allow_inconsistencies) {
for (const std::string& expected_output : output_nodes) {
if (!new_node_names.count(expected_output)) {
LOG(WARNING) << "Expected " << expected_output
<< " to be preserved.";
abort_replacement = true;
}
}
}
if (abort_replacement) {
LOG(WARNING) << "Generator function didn't preserve needed nodes, "
<< "copying old replacements back in instead.";
std::vector<NodeDef> old_nodes;
MatchedNodesAsArray(*match, &old_nodes);
for (const NodeDef& old_node : old_nodes) {
NodeDef* added_node = output_graph_def->mutable_node()->Add();
*added_node = old_node;
}
} else {
for (const NodeDef& new_node : new_nodes) {
NodeDef* added_node = output_graph_def->mutable_node()->Add();
*added_node = new_node;
}
}
} else if (!matched_nodes.count(input_node.name())) {
// This node isn't part of any match, so just copy it over.
NodeDef* added_node = output_graph_def->mutable_node()->Add();
*added_node = input_node;
} else {
// Do nothing, because this is an internal part of a matching subgraph,
// and so will have been replaced by a new replacement subgraph.
}
}
return absl::OkStatus();
}
absl::Status RenameNodeInputs(
const GraphDef& input_graph_def,
const std::map<std::string, std::string>& inputs_to_rename,
const std::unordered_set<std::string>& nodes_to_ignore,
GraphDef* output_graph_def) {
std::map<std::string, std::vector<std::pair<std::string, std::string>>>
canonical_inputs_to_rename;
for (const auto& input_to_rename : inputs_to_rename) {
canonical_inputs_to_rename[NodeNameFromInput(input_to_rename.first)]
.push_back({input_to_rename.first, input_to_rename.second});
}
output_graph_def->Clear();
for (const NodeDef& node : input_graph_def.node()) {
NodeDef* new_node = output_graph_def->mutable_node()->Add();
*new_node = node;
new_node->mutable_input()->Clear();
for (const std::string& input_name : node.input()) {
std::set<std::string> already_visited;
std::string new_input_name = input_name;
while (
canonical_inputs_to_rename.count(NodeNameFromInput(new_input_name))) {
std::string input_node_name = NodeNameFromInput(new_input_name);
if (already_visited.count(input_node_name)) {
return absl::InvalidArgumentError(
absl::StrCat("RenameNodeInputs argument contains a cycle for ",
input_node_name));
}
already_visited.insert(input_node_name);
if (nodes_to_ignore.count(node.name())) {
break;
}
bool any_match_found = false;
for (const std::pair<std::string, std::string>& input_to_rename :
canonical_inputs_to_rename.at(input_node_name)) {
const std::string& source_name = input_to_rename.first;
const std::string& dest_name = input_to_rename.second;
bool is_match;
std::string match_name;
if (absl::EndsWith(source_name, ":*")) {
is_match = true;
std::string prefix;
std::string unused_node_name;
std::string suffix;
NodeNamePartsFromInput(new_input_name, &prefix, &unused_node_name,
&suffix);
match_name = prefix + dest_name + suffix;
} else {
is_match = (CanonicalInputName(source_name) ==
CanonicalInputName(new_input_name));
match_name = dest_name;
}
if (is_match) {
new_input_name = match_name;
any_match_found = true;
}
}
if (!any_match_found) {
break;
}
}
*(new_node->mutable_input()->Add()) = new_input_name;
}
}
return absl::OkStatus();
}
void CopyOriginalMatch(const NodeMatch& match,
std::vector<NodeDef>* new_nodes) {
std::vector<NodeDef> old_nodes;
MatchedNodesAsArray(match, &old_nodes);
for (const NodeDef& old_node : old_nodes) {
new_nodes->push_back(old_node);
}
}
TransformRegistry* GetTransformRegistry() {
static TransformRegistry transform_registry;
return &transform_registry;
}
void FindInvalidInputs(
const GraphDef& graph_def,
std::vector<std::pair<std::string, std::string>>* invalid_inputs) {
std::map<std::string, const NodeDef*> node_map;
MapNamesToNodes(graph_def, &node_map);
for (const NodeDef& node : graph_def.node()) {
for (const std::string& input : node.input()) {
std::string input_node = NodeNameFromInput(input);
if (!node_map.count(input_node)) {
invalid_inputs->push_back({node.name(), input_node});
}
}
}
}
absl::Status IsGraphValid(const GraphDef& graph_def) {
std::vector<std::pair<std::string, std::string>> invalid_inputs;
FindInvalidInputs(graph_def, &invalid_inputs);
if (!invalid_inputs.empty()) {
std::map<std::string, const NodeDef*> node_map;
MapNamesToNodes(graph_def, &node_map);
for (const std::pair<std::string, std::string>& invalid_input :
invalid_inputs) {
LOG(ERROR) << "Invalid input " << invalid_input.second << " for node "
<< invalid_input.first << " - "
<< node_map[invalid_input.first]->DebugString();
}
return absl::InternalError(
"Invalid graph with inputs referring to nonexistent nodes");
}
return absl::OkStatus();
}
absl::Status GetInOutTypes(const NodeDef& node_def, DataTypeVector* inputs,
DataTypeVector* outputs) {
const OpDef* op_def;
TF_RETURN_IF_ERROR(OpRegistry::Global()->LookUpOpDef(node_def.op(), &op_def));
TF_RETURN_IF_ERROR(InOutTypesForNode(node_def, *op_def, inputs, outputs));
return absl::OkStatus();
}
absl::Status TensorShapeFromString(const std::string& shape_string,
TensorShape* result) {
if (shape_string.empty()) {
return absl::InvalidArgumentError("Specified shape is empty.");
}
std::vector<std::string> dims_as_str = str_util::Split(shape_string, ",");
std::vector<int64_t> dims;
for (const std::string& dim : dims_as_str) {
int64_t tmp;
if (absl::SimpleAtoi(dim, &tmp)) {
dims.push_back(tmp);
} else {
return absl::InvalidArgumentError(
absl::StrCat("Could parse as shape: '", shape_string, "'"));
}
}
*result = TensorShape(dims);
return absl::OkStatus();
}
int TransformFuncContext::CountParameters(const std::string& name) const {
if (params.count(name)) {
return params.at(name).size();
} else {
return 0;
}
}
absl::Status TransformFuncContext::GetOneStringParameter(
const std::string& name, const std::string& default_value,
std::string* result) const {
const int params_count = CountParameters(name);
if (params_count == 0) {
*result = default_value;
return absl::OkStatus();
} else if (params_count == 1) {
*result = params.at(name).at(0);
return absl::OkStatus();
} else {
return absl::InvalidArgumentError(
absl::StrCat("Expected a single '", name, "' parameter, but found ",
params_count, " occurrences"));
}
}
absl::Status TransformFuncContext::GetOneInt32Parameter(const std::string& name,
int32_t default_value,
int32_t* result) const {
const int params_count = CountParameters(name);
if (params_count == 0) {
*result = default_value;
return absl::OkStatus();
}
std::string string_value;
TF_RETURN_IF_ERROR(GetOneStringParameter(name, "", &string_value));
if (!absl::SimpleAtoi(absl::string_view(string_value), result)) {
return absl::InvalidArgumentError(
absl::StrCat("Couldn't interpret the ", name,
" argument as a number:", string_value));
}
return absl::OkStatus();
}
absl::Status TransformFuncContext::GetOneInt64Parameter(const std::string& name,
int64_t default_value,
int64_t* result) const {
const int params_count = CountParameters(name);
if (params_count == 0) {
*result = default_value;
return absl::OkStatus();
}
std::string string_value;
TF_RETURN_IF_ERROR(GetOneStringParameter(name, "", &string_value));
if (!absl::SimpleAtoi(absl::string_view(string_value), result)) {
return absl::InvalidArgumentError(
absl::StrCat("Couldn't interpret the ", name,
" argument as a number:", string_value));
}
return absl::OkStatus();
}
absl::Status TransformFuncContext::GetOneFloatParameter(const std::string& name,
float default_value,
float* result) const {
const int params_count = CountParameters(name);
if (params_count == 0) {
*result = default_value;
return absl::OkStatus();
}
std::string string_value;
TF_RETURN_IF_ERROR(GetOneStringParameter(name, "", &string_value));
if (!absl::SimpleAtof(string_value.c_str(), result)) {
return absl::InvalidArgumentError(
absl::StrCat("Couldn't interpret the ", name,
" argument as a float number:", string_value));
}
return absl::OkStatus();
}
absl::Status TransformFuncContext::GetOneBoolParameter(const std::string& name,
bool default_value,
bool* result) const {
const int params_count = CountParameters(name);
if (params_count == 0) {
*result = default_value;
return absl::OkStatus();
}
std::string string_value;
TF_RETURN_IF_ERROR(GetOneStringParameter(name, "", &string_value));
if (string_value == "true" || string_value == "1") {
*result = true;
} else if (string_value == "false" || string_value == "0") {
*result = false;
} else {
return absl::InvalidArgumentError(
absl::StrCat("Couldn't interpret the ", name, " argument as a boolean:",
string_value, " (expected true, false, 0 or 1)"));
}
return absl::OkStatus();
}
} // namespace graph_transforms
} // namespace tensorflow
tensorflow/tools/graph_transforms/transform_utils.h
+298
View File
@@ -0,0 +1,298 @@
/* Copyright 2015 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_TOOLS_GRAPH_TRANSFORMS_TRANSFORM_UTILS_H_
#define TENSORFLOW_TOOLS_GRAPH_TRANSFORMS_TRANSFORM_UTILS_H_
#include <set>
#include <unordered_set>
#include <vector>
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/attr_value_util.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor.pb.h"
#include "tensorflow/core/lib/core/status.h"
namespace tensorflow {
namespace graph_transforms {
// Used to quickly look up nodes in the graph def from a name.
void MapNamesToNodes(const GraphDef& graph_def,
std::map<std::string, const NodeDef*>* result);
// For every node in the graph create a list of the nodes that use it as an
// input.
void MapNodesToOutputs(
const GraphDef& graph_def,
std::map<std::string, std::vector<const NodeDef*>>* result);
// NodeDef input strings can contain other information besides the name of an
// input node. These include:
// - Optional '^' prefix, indicating this is a control edge.
// - The required name of the input node.
// - Optional ':<number>' suffix, showing which output of the node to use.
// This function takes a raw string, and breaks it into those component parts.
// The rules for inputs in function libraries are a bit more complex, and
// aren't handled by this routine.
void NodeNamePartsFromInput(const std::string& input_name, std::string* prefix,
std::string* node_name, std::string* suffix);
// Adds a ':0' port to any inputs with no suffix, to make comparisons easier.
std::string CanonicalInputName(const std::string& input_name);
// Convenience function to strip the optional prefix and suffix components from
// a string pulled from a NodeDef input, and return the plain node name.
std::string NodeNameFromInput(const std::string& input_name);
// Returns a stable hash for the contents of the NodeDef, so that equivalent
// nodes should have equal hashes.
uint64_t HashNodeDef(const NodeDef& node);
// Adds the given node name to the end of the node's inputs.
void AddNodeInput(const std::string& input_name, NodeDef* node);
// Copies an attribute from one NodeDef to another.
void CopyNodeAttr(const NodeDef& source, const std::string& source_key,
const std::string& dest_key, NodeDef* dest);
// Inserts a value into a NodeDef's map of attributes.
// This is a bit different than AddNodeAttr in node_def_util.h because it
// overwrites any existing attributes with the same key.
template <class T>
inline void SetNodeAttr(const std::string& key, const T& value, NodeDef* node) {
AttrValue attr_value;
SetAttrValue(value, &attr_value);
auto* attr_map = node->mutable_attr();
(*attr_map)[key] = attr_value;
}
template <class T>
inline void SetNodeTensorAttr(const std::string& key, const Tensor& tensor,
NodeDef* node) {
TensorProto tensor_proto;
tensor.AsProtoTensorContent(&tensor_proto);
SetNodeAttr(key, tensor_proto, node);
}
// Inserts a Tensor into the specified attribute of a NodeDef.
template <class T>
inline void SetNodeTensorAttr(const std::string& key, const TensorShape& shape,
const std::vector<T>& values, NodeDef* node) {
const DataType dtype = DataTypeToEnum<T>::v();
CHECK_EQ(shape.num_elements(), values.size());
Tensor tensor(dtype, shape);
T* dest_data = tensor.flat<T>().data();
std::copy_n(values.data(), values.size(), dest_data);
SetNodeTensorAttr<T>(key, tensor, node);
}
// Retrieves a tensor value from a NodeDef attribute.
Tensor GetNodeTensorAttr(const NodeDef& node, const std::string& key);
// Creates a copy of the input GraphDef, but only containing the nodes where the
// supplied selector function returned true.
void FilterGraphDef(const GraphDef& input_graph_def,
std::function<bool(const NodeDef&)> selector,
GraphDef* output_graph_def);
// Creates a copy of the input graph, with all occurrences of the attributes
// with the names in the argument removed from the node defs.
void RemoveAttributes(const GraphDef& input_graph_def,
const std::vector<std::string>& attributes,
GraphDef* output_graph_def);
// For a lot of replacement and matching operations it's useful to have the
// nodes processed in a controlled order, so this does a topological sort to
// ensure that nodes always appear in the GraphDef.node list after their inputs.
absl::Status SortByExecutionOrder(const GraphDef& input_graph_def,
GraphDef* output_graph_def);
// Finds inputs that refer to nodes that are not in the graph.
void FindInvalidInputs(
const GraphDef& graph_def,
std::vector<std::pair<std::string, std::string>>* invalid_inputs);
// Returns a descriptive error status if there are problems spotted with the
// graph.
absl::Status IsGraphValid(const GraphDef& graph_def);
// Returns input and output types for a particular NodeDef.
absl::Status GetInOutTypes(const NodeDef& node_def, DataTypeVector* inputs,
DataTypeVector* outputs);
// Takes a comma-separated string of numbers and parses them into a shape.
absl::Status TensorShapeFromString(const std::string& shape_string,
TensorShape* result);
// This is used to spot particular subgraphs in a larger model. To use it,
// create a pattern like:
// OpTypePattern pattern({"Conv2D", {{"ResizeBilinear", {{"MirrorPad"}}}}});
// This defines a subgraph where a Conv2D has a ResizeBilinear input, which
// pulls from a MirrorPad op.
// Regular expressions aren't supported for the op names, but you can use "*" to
// match any op. You can also use | as a separator to match multiple op names,
// like "Reshape|Concat|Conv2D".
struct OpTypePattern {
std::string op;
std::vector<OpTypePattern> inputs;
std::string DebugString() const;
};
// Returns a sub-graph of nodes that match a pattern.
struct NodeMatch {
NodeMatch() : node() {}
NodeDef node;
std::vector<NodeMatch> inputs;
std::string DebugString() const;
};
// Utility class to spot subgraphs matching particular patterns.
class GraphMatcher {
public:
GraphMatcher(const GraphDef& graph_def);
// Sorts the input nodes into execution order, and then skips any previously
// matches so that no node appears in more than one match. The NodeDef
// pointers contained in the results are owned by the GraphMatcher object, and
// so will be invalid after its lifetime.
absl::Status GetOpTypeMatches(const OpTypePattern& pattern,
std::vector<NodeMatch>* matches);
private:
bool DoesOpTypeMatch(const NodeDef& node, const OpTypePattern& pattern,
const std::set<std::string>& previously_matched_nodes,
NodeMatch* match);
GraphDef graph_def_;
std::map<std::string, const NodeDef*> node_map_;
};
struct ReplaceMatchingOpTypesOptions {
// Whether to raise an error if the graph is left with dangling inputs. If you
// enable this option, you must fix inconsistencies in a later pass.
bool allow_inconsistencies;
};
// Replaces all of the matching sub-graphs with new ops. This calls into the
// given function, and expects to receive a set of new nodes to replace each
// matched sub-graph. It has some logic to protect the integrity of the
// resulting graph, for example making sure that nodes needed by other nodes
// outside the sub-graph aren't removed. These are passed in as the set of
// outputs, and nodes with the same names must be added to the new nodes
// produced by the replacement function. Many of these checks can be disabled
// by setting allow_inconsistencies to true in the options, but then it's the
// caller's responsibility to patch up any problems before passing on the graph
// to others. There's more comprehensive usage documentation in the README.
absl::Status ReplaceMatchingOpTypes(
const GraphDef& input_graph_def, const OpTypePattern& pattern,
const std::function<absl::Status(
const NodeMatch&, const std::set<std::string>&,
const std::set<std::string>&, std::vector<NodeDef>*)>& node_generator,
const ReplaceMatchingOpTypesOptions& options, GraphDef* output_graph_def);
// Returns a list of the unique nodes found in this match.
void MatchedNodesAsArray(const NodeMatch& match, std::vector<NodeDef>* result);
// Changes all input references to a particular node name. Any nodes with names
// listed in nodes_to_ignore will not have their inputs rewritten.
absl::Status RenameNodeInputs(
const GraphDef& input_graph_def,
const std::map<std::string, std::string>& inputs_to_rename,
const std::unordered_set<std::string>& nodes_to_ignore,
GraphDef* output_graph_def);
// Utility function that copies all the nodes found in a match into the
// new_nodes list. This is useful in replacement functions when you decide to
// leave the original matched subgraph untouched and make no changes.
void CopyOriginalMatch(const NodeMatch& match, std::vector<NodeDef>* new_nodes);
// Holds information that's needed for transform functions.
typedef std::map<std::string, std::vector<std::string>> TransformFuncParameters;
struct TransformFuncContext {
std::vector<std::string> input_names;
std::vector<std::string> output_names;
TransformFuncParameters params;
// Returns how many occurrences of the given parameter are present.
int CountParameters(const std::string& name) const;
// Gets a single instance of a parameter, using a default if it's not present.
absl::Status GetOneStringParameter(const std::string& name,
const std::string& default_value,
std::string* result) const;
// Gets a single occurrence of a parameter as a 32-bit integer, falling back
// to a default if it isn't present and returning an error if it isn't
// convertible to a number.
absl::Status GetOneInt32Parameter(const std::string& name,
int32_t default_value,
int32_t* result) const;
// Gets a single occurrence of a parameter as a 64-bit integer, falling back
// to a default if it isn't present and returning an error if it isn't
// convertible to a number.
absl::Status GetOneInt64Parameter(const std::string& name,
int64_t default_value,
int64_t* result) const;
// Gets a single occurrence of a parameter as a floating point number, falling
// back to a default if it isn't present and returning an error if it isn't
// convertible to a number.
absl::Status GetOneFloatParameter(const std::string& name,
float default_value, float* result) const;
// Gets a single occurrence of a parameter as a boolean, falling back to a
// default if it isn't present and returning an error if it's not one of
// "true", "1", "false", or "0".
absl::Status GetOneBoolParameter(const std::string& name, bool default_value,
bool* result) const;
};
// This is the function API for all graph transformations, taking an input
// GraphDef and other arguments, and returning a transformed GraphDef.
typedef std::function<absl::Status(
const GraphDef&, const TransformFuncContext& context, GraphDef*)>
TransformFunc;
// To add a new graph transform function, call the macro:
// REGISTER_GRAPH_TRANSFORM("fold_constants", FoldConstants);
// Under the hood this adds the function to the list of known transforms, so you
// just need to link in the .cc file with your registration call to have access
// to it through the command line tool.
// The rest of the machinery below is to enable that automagical registration.
typedef std::map<std::string, TransformFunc> TransformRegistry;
TransformRegistry* GetTransformRegistry();
class TransformRegistrar {
public:
TransformRegistrar(const std::string& name, TransformFunc transform_func) {
TransformRegistry* transform_registry = GetTransformRegistry();
(*transform_registry)[name] = transform_func;
}
};
#define REGISTER_GRAPH_TRANSFORM(name, func) \
REGISTER_GRAPH_TRANSFORM_UNIQ_HELPER(__COUNTER__, name, func)
#define REGISTER_GRAPH_TRANSFORM_UNIQ_HELPER(ctr, name, func) \
REGISTER_GRAPH_TRANSFORM_UNIQ(ctr, name, func)
#define REGISTER_GRAPH_TRANSFORM_UNIQ(ctr, name, func) \
static tensorflow::graph_transforms::TransformRegistrar \
registrar__body__##ctr##__object(name, func);
} // namespace graph_transforms
} // namespace tensorflow
#endif // TENSORFLOW_TOOLS_GRAPH_TRANSFORMS_TRANSFORM_UTILS_H_
tensorflow/tools/graph_transforms/transform_utils_test.cc
+1165
View File
File diff suppressed because it is too large Load Diff