1# pylint: disable=g-bad-file-header
2# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#     http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15# ==============================================================================
16"""Utilities to remove unneeded nodes from a GraphDefs."""
17
18from __future__ import absolute_import
19from __future__ import division
20from __future__ import print_function
21import copy
22
23from google.protobuf import text_format
24
25from tensorflow.core.framework import attr_value_pb2
26from tensorflow.core.framework import graph_pb2
27from tensorflow.core.framework import node_def_pb2
28from tensorflow.python.framework import graph_util
29from tensorflow.python.platform import gfile
30
31
32def strip_unused(input_graph_def, input_node_names, output_node_names,
33                 placeholder_type_enum):
34  """Removes unused nodes from a GraphDef.
35
36  Args:
37    input_graph_def: A graph with nodes we want to prune.
38    input_node_names: A list of the nodes we use as inputs.
39    output_node_names: A list of the output nodes.
40    placeholder_type_enum: The AttrValue enum for the placeholder data type, or
41        a list that specifies one value per input node name.
42
43  Returns:
44    A `GraphDef` with all unnecessary ops removed.
45
46  Raises:
47    ValueError: If any element in `input_node_names` refers to a tensor instead
48      of an operation.
49    KeyError: If any element in `input_node_names` is not found in the graph.
50  """
51  for name in input_node_names:
52    if ":" in name:
53      raise ValueError("Name '%s' appears to refer to a Tensor, "
54                       "not a Operation." % name)
55
56  # Here we replace the nodes we're going to override as inputs with
57  # placeholders so that any unused nodes that are inputs to them are
58  # automatically stripped out by extract_sub_graph().
59  not_found = {name for name in input_node_names}
60  inputs_replaced_graph_def = graph_pb2.GraphDef()
61  for node in input_graph_def.node:
62    if node.name in input_node_names:
63      not_found.remove(node.name)
64      placeholder_node = node_def_pb2.NodeDef()
65      placeholder_node.op = "Placeholder"
66      placeholder_node.name = node.name
67      if isinstance(placeholder_type_enum, list):
68        input_node_index = input_node_names.index(node.name)
69        placeholder_node.attr["dtype"].CopyFrom(
70            attr_value_pb2.AttrValue(type=placeholder_type_enum[
71                input_node_index]))
72      else:
73        placeholder_node.attr["dtype"].CopyFrom(
74            attr_value_pb2.AttrValue(type=placeholder_type_enum))
75      if "_output_shapes" in node.attr:
76        placeholder_node.attr["_output_shapes"].CopyFrom(node.attr[
77            "_output_shapes"])
78      inputs_replaced_graph_def.node.extend([placeholder_node])
79    else:
80      inputs_replaced_graph_def.node.extend([copy.deepcopy(node)])
81
82  if not_found:
83    raise KeyError("The following input nodes were not found: %s\n" % not_found)
84
85  output_graph_def = graph_util.extract_sub_graph(inputs_replaced_graph_def,
86                                                  output_node_names)
87  return output_graph_def
88
89
90def strip_unused_from_files(input_graph, input_binary, output_graph,
91                            output_binary, input_node_names, output_node_names,
92                            placeholder_type_enum):
93  """Removes unused nodes from a graph file."""
94
95  if not gfile.Exists(input_graph):
96    print("Input graph file '" + input_graph + "' does not exist!")
97    return -1
98
99  if not output_node_names:
100    print("You need to supply the name of a node to --output_node_names.")
101    return -1
102
103  input_graph_def = graph_pb2.GraphDef()
104  mode = "rb" if input_binary else "r"
105  with gfile.GFile(input_graph, mode) as f:
106    if input_binary:
107      input_graph_def.ParseFromString(f.read())
108    else:
109      text_format.Merge(f.read(), input_graph_def)
110
111  output_graph_def = strip_unused(input_graph_def,
112                                  input_node_names.split(","),
113                                  output_node_names.split(","),
114                                  placeholder_type_enum)
115
116  if output_binary:
117    with gfile.GFile(output_graph, "wb") as f:
118      f.write(output_graph_def.SerializeToString())
119  else:
120    with gfile.GFile(output_graph, "w") as f:
121      f.write(text_format.MessageToString(output_graph_def))
122  print("%d ops in the final graph." % len(output_graph_def.node))
123