Home | History | Annotate | Download | only in graph_transforms 
      1 /* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
      2 
      3 Licensed under the Apache License, Version 2.0 (the "License");
      4 you may not use this file except in compliance with the License.
      5 You may obtain a copy of the License at
      6 
      7     http://www.apache.org/licenses/LICENSE-2.0
      8 
      9 Unless required by applicable law or agreed to in writing, software
     10 distributed under the License is distributed on an "AS IS" BASIS,
     11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
     12 See the License for the specific language governing permissions and
     13 limitations under the License.
     14 ==============================================================================*/
     15 
     16 #include "tensorflow/tools/graph_transforms/fold_constants_lib.h"
     17 
     18 #include "tensorflow/core/common_runtime/constant_folding.h"
     19 #include "tensorflow/core/graph/graph_constructor.h"
     20 #include "tensorflow/core/graph/node_builder.h"
     21 #include "tensorflow/core/graph/subgraph.h"
     22 #include "tensorflow/core/platform/init_main.h"
     23 #include "tensorflow/core/public/session.h"
     24 #include "tensorflow/core/util/command_line_flags.h"
     25 #include "tensorflow/tools/graph_transforms/transform_utils.h"
     26 
     27 namespace tensorflow {
     28 namespace graph_transforms {
     29 
     30 // Deletes any specified types of nodes, unless they're necessary for the
     31 // graph's inputs or outputs.
     32 Status RemoveNodes(const GraphDef& input_graph_def,
     33                    const TransformFuncContext& context,
     34                    GraphDef* output_graph_def) {
     35   if (!context.params.count("op")) {
     36     return errors::InvalidArgument(
     37         "remove_nodes expects at least one 'op'"
     38         "argument, e.g. remove_nodes(op=Identity)");
     39   }
     40   int32 max_inputs;
     41   TF_RETURN_IF_ERROR(
     42       context.GetOneInt32Parameter("max_inputs", 1, &max_inputs));
     43 
     44   // Make sure we don't get rid of any nodes used as graph inputs or outputs.
     45   std::set<string> required_nodes;
     46   for (const string& input : context.input_names) {
     47     required_nodes.insert(NodeNameFromInput(input));
     48   }
     49   for (const string& output : context.output_names) {
     50     required_nodes.insert(NodeNameFromInput(output));
     51   }
     52 
     53   std::vector<string> ops_to_remove = context.params.at("op");
     54   GraphDef current_graph_def = input_graph_def;
     55   for (const string& op : ops_to_remove) {
     56     for (int num_inputs = 1; num_inputs <= max_inputs; ++num_inputs) {
     57       // Look for a variable number of inputs.
     58       OpTypePattern pattern = {op};
     59       pattern.inputs.resize(num_inputs);
     60       for (int i = 0; i < num_inputs; ++i) {
     61         pattern.inputs[i] = {"*"};
     62       }
     63       // Keep looking for nodes to remove until there are no more changes.
     64       bool any_nodes_removed;
     65       do {
     66         any_nodes_removed = false;
     67         std::map<string, string> inputs_to_rename;
     68         GraphDef replaced_graph_def;
     69         TF_RETURN_IF_ERROR(ReplaceMatchingOpTypes(
     70             current_graph_def, pattern,
     71             [&inputs_to_rename, &required_nodes, &any_nodes_removed](
     72                 const NodeMatch& match, const std::set<string>& input_nodes,
     73                 const std::set<string>& output_nodes,
     74                 std::vector<NodeDef>* new_nodes) {
     75               const NodeDef& replace_node = match.node;
     76               // If this node is needed in the inputs or outputs don't replace
     77               // it.
     78               if (required_nodes.count(replace_node.name())) {
     79                 LOG(INFO) << "Skipping replacement for " << replace_node.name();
     80                 CopyOriginalMatch(match, new_nodes);
     81                 return Status::OK();
     82               }
     83               const NodeDef& input_node = match.inputs[0].node;
     84               string target_name = input_node.name();
     85               for (const string& input : replace_node.input()) {
     86                 if (!input.compare(0, target_name.size(), target_name)) {
     87                   if (input.size() == target_name.size() ||
     88                       input[target_name.size()] == ':') {
     89                     target_name = input;
     90                     break;
     91                   }
     92                 }
     93               }
     94               inputs_to_rename[replace_node.name()] = target_name;
     95               inputs_to_rename["^" + replace_node.name()] =
     96                   "^" + input_node.name();
     97               new_nodes->push_back(input_node);
     98               any_nodes_removed = true;
     99               return Status::OK();
    100             },
    101             {true}, &replaced_graph_def));
    102         // Make sure all references to removed nodes now point to their inputs.
    103         TF_RETURN_IF_ERROR(
    104             RenameNodeInputs(replaced_graph_def, inputs_to_rename,
    105                              std::unordered_set<string>(), &current_graph_def));
    106       } while (any_nodes_removed);
    107     }
    108   }
    109 
    110   *output_graph_def = current_graph_def;
    111   return Status::OK();
    112 }
    113 
    114 REGISTER_GRAPH_TRANSFORM("remove_nodes", RemoveNodes);
    115 
    116 }  // namespace graph_transforms
    117 }  // namespace tensorflow
    118