Tool to fold constants in saved graphs.

It's often useful to preprocess graphs to fold constant expression subgraphs
into single constant nodes, especially as preparation for further processing for
purposes like batch normalization folding. This change uses the built-in constant
folding in the runtime to rewrite frozen models on disk.
Change: 135440279

tensorflow/tools/graph_transforms/BUILD

@@ -0,0 +1,106 @@
+# Description:
+#   Utilities that perform useful transformations on graphs
+
+package(default_visibility = ["//visibility:public"])
+
+licenses(["notice"])  # Apache 2.0
+
+load(
+    "//tensorflow:tensorflow.bzl",
+    "tf_copts",
+    "tf_cc_test",
+)
+
+exports_files(["LICENSE"])
+
+cc_library(
+    name = "transform_utils",
+    srcs = [
+        "transform_utils.cc",
+    ],
+    hdrs = [
+        "transform_utils.h",
+    ],
+    copts = tf_copts(),
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tensorflow/core:core_cpu",
+        "//tensorflow/core:core_cpu_internal",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:framework_internal",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:tensorflow",
+    ],
+)
+
+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",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core:testlib",
+    ],
+)
+
+cc_library(
+    name = "fold_constants_lib",
+    srcs = [
+        "fold_constants_lib.cc",
+    ],
+    hdrs = [
+        "fold_constants_lib.h",
+    ],
+    copts = tf_copts(),
+    visibility = ["//visibility:public"],
+    deps = [
+        ":transform_utils",
+        "//tensorflow/core:core_cpu",
+        "//tensorflow/core:core_cpu_internal",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:framework_internal",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:tensorflow",
+    ],
+)
+
+tf_cc_test(
+    name = "fold_constants_test",
+    size = "small",
+    srcs = ["fold_constants_test.cc"],
+    deps = [
+        ":fold_constants_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",
+    ],
+)
+
+cc_binary(
+    name = "fold_constants_tool",
+    srcs = ["fold_constants_tool.cc"],
+    copts = tf_copts(),
+    linkstatic = 1,
+    visibility = ["//visibility:public"],
+    deps = [
+        ":fold_constants_lib",
+        "//tensorflow/core:framework_internal",
+        "//tensorflow/core:lib",
+    ],
+)

tensorflow/tools/graph_transforms/fold_constants_lib.cc

@@ -0,0 +1,159 @@
+/* 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/fold_constants_lib.h"
+
+#include "tensorflow/core/common_runtime/constant_folding.h"
+#include "tensorflow/core/graph/graph_constructor.h"
+#include "tensorflow/core/graph/node_builder.h"
+#include "tensorflow/core/graph/subgraph.h"
+#include "tensorflow/core/public/session.h"
+#include "tensorflow/tools/graph_transforms/transform_utils.h"
+
+namespace tensorflow {
+namespace graph_transforms {
+
+Status ReplaceSendRecvs(const GraphDef& original_graph_def,
+                        const GraphDef& rewritten_graph_def,
+                        const std::vector<string>& inputs,
+                        const std::vector<string>& outputs,
+                        GraphDef* output_graph_def) {
+  std::map<string, const NodeDef*> original_map;
+  MapNamesToNodes(original_graph_def, &original_map);
+  std::map<string, string> new_node_names;
+  for (const NodeDef& node : rewritten_graph_def.node()) {
+    // If the op isn't a Recv, or it was in the original, nothing to do.
+    if ((node.op() != "_Recv") || (original_map.count(node.name()) == 1)) {
+      continue;
+    }
+    // See if it matches an input from the original.
+    for (const string& input : inputs) {
+      // Here we rely on the naming convention for the Recv nodes that
+      // RewriteGraphForExecution adds in the place of the feed inputs.
+      string input_prefix = "_recv_" + input + "_";
+      if (StringPiece(node.name()).starts_with(input_prefix)) {
+        // If it does, prepare to rename any inputs that refer to it.
+        new_node_names[node.name()] = input;
+      }
+    }
+  }
+
+  std::vector<NodeDef> nodes_to_add;
+  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;
+    new_node.CopyFrom(node);
+    new_node.mutable_input()->Clear();
+    for (const string& old_input : node.input()) {
+      string input_prefix;
+      string input_node_name;
+      string input_suffix;
+      NodeNamePartsFromInput(old_input, &input_prefix, &input_node_name,
+                             &input_suffix);
+      string new_input;
+      if (new_node_names.count(input_node_name) > 0) {
+        new_input =
+            input_prefix + new_node_names[input_node_name] + input_suffix;
+      } else {
+        new_input = old_input;
+      }
+      *(new_node.mutable_input()->Add()) = new_input;
+    }
+    nodes_to_add.push_back(new_node);
+  }
+  for (std::pair<string, string> entry : new_node_names) {
+    string removed_node_name = entry.second;
+    const NodeDef* removed_node = original_map[removed_node_name];
+    NodeDef new_node;
+    new_node.CopyFrom(*removed_node);
+    nodes_to_add.push_back(new_node);
+  }
+
+  for (const NodeDef& node : nodes_to_add) {
+    output_graph_def->mutable_node()->Add()->CopyFrom(node);
+  }
+  return Status::OK();
+}
+
+Status RemoveUnusedNodes(const GraphDef& input_graph_def,
+                         const std::vector<string>& inputs,
+                         const std::vector<string>& outputs,
+                         GraphDef* output_graph_def) {
+  std::map<string, const NodeDef*> node_map;
+  MapNamesToNodes(input_graph_def, &node_map);
+
+  std::map<string, bool> used_nodes;
+  for (const string& input : inputs) {
+    used_nodes[input] = true;
+  }
+  std::vector<string> current_nodes = outputs;
+  while (!current_nodes.empty()) {
+    std::vector<string> next_nodes;
+    for (const string& node_name : current_nodes) {
+      used_nodes[node_name] = true;
+      if (node_map.count(node_name) == 0) {
+        LOG(ERROR) << "Bad graph structure, no node named '" << node_name
+                   << "' found for input lookup";
+        return errors::InvalidArgument("Bad graph structure, no node named '",
+                                       node_name, "' found for input lookup");
+      }
+      const NodeDef& node = *(node_map[node_name]);
+      for (const string& input_name : node.input()) {
+        const string& input_node_name = NodeNameFromInput(input_name);
+        if (used_nodes.count(input_node_name) == 0) {
+          next_nodes.push_back(input_node_name);
+        }
+      }
+    }
+    current_nodes = next_nodes;
+  }
+  FilterGraphDef(
+      input_graph_def,
+      [&](const NodeDef& node) { return used_nodes.count(node.name()) > 0; },
+      output_graph_def);
+
+  return Status::OK();
+}
+
+Status FoldConstants(const GraphDef& input_graph_def,
+                     const std::vector<string>& inputs,
+                     const std::vector<string>& outputs,
+                     GraphDef* output_graph_def) {
+  Graph input_graph(OpRegistry::Global());
+  ImportGraphDefOptions import_opts;
+  TF_RETURN_IF_ERROR(
+      ImportGraphDef(import_opts, input_graph_def, &input_graph, nullptr));
+  DeviceAttributes device_attributes;
+  TF_RETURN_IF_ERROR(subgraph::RewriteGraphForExecution(
+      &input_graph, inputs, outputs, {}, device_attributes));
+  DoConstantFolding(ConstantFoldingOptions(), nullptr, Env::Default(), nullptr,
+                    &input_graph);
+  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, inputs,
+                                      outputs, &send_recvs_replaced));
+  TF_RETURN_IF_ERROR(RemoveUnusedNodes(send_recvs_replaced, inputs, outputs,
+                                       output_graph_def));
+  return Status::OK();
+}
+
+}  // namespace graph_transforms
+}  // namespace tensorflow

tensorflow/tools/graph_transforms/fold_constants_lib.h

@@ -0,0 +1,44 @@
+/* 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_H_
+#define TENSORFLOW_TOOLS_GRAPH_TRANSFORMS_FOLD_CONSTANTS_H_
+
+#include "tensorflow/core/framework/graph.pb.h"
+#include "tensorflow/core/lib/core/status.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.
+Status FoldConstants(const GraphDef& input_graph_def,
+                     const std::vector<string>& inputs,
+                     const std::vector<string>& outputs,
+                     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.
+Status RemoveUnusedNodes(const GraphDef& input_graph_def,
+                         const std::vector<string>& inputs,
+                         const std::vector<string>& outputs,
+                         GraphDef* output_graph_def);
+
+}  // namespace graph_transforms
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_TOOLS_GRAPH_TRANSFORMS_FOLD_CONSTANTS_H_