Deletes any specified types of nodes, unless they're necessary for the graph's inputs or outputs.
| 29 | // Deletes any specified types of nodes, unless they're necessary for the |
| 30 | // graph's inputs or outputs. |
| 31 | Status RemoveNodes(const GraphDef& input_graph_def, |
| 32 | const TransformFuncContext& context, |
| 33 | GraphDef* output_graph_def) { |
| 34 | if (!context.params.count("op")) { |
| 35 | return errors::InvalidArgument( |
| 36 | "remove_nodes expects at least one 'op'" |
| 37 | "argument, e.g. remove_nodes(op=Identity)"); |
| 38 | } |
| 39 | int32 max_inputs; |
| 40 | TF_RETURN_IF_ERROR( |
| 41 | context.GetOneInt32Parameter("max_inputs", 1, &max_inputs)); |
| 42 | |
| 43 | // Make sure we don't get rid of any nodes used as graph inputs or outputs. |
| 44 | std::set<string> required_nodes; |
| 45 | for (const string& input : context.input_names) { |
| 46 | required_nodes.insert(NodeNameFromInput(input)); |
| 47 | } |
| 48 | for (const string& output : context.output_names) { |
| 49 | required_nodes.insert(NodeNameFromInput(output)); |
| 50 | } |
| 51 | |
| 52 | std::vector<string> ops_to_remove = context.params.at("op"); |
| 53 | GraphDef current_graph_def = input_graph_def; |
| 54 | for (const string& op : ops_to_remove) { |
| 55 | for (int num_inputs = 1; num_inputs <= max_inputs; ++num_inputs) { |
| 56 | // Look for a variable number of inputs. |
| 57 | OpTypePattern pattern = {op}; |
| 58 | pattern.inputs.resize(num_inputs); |
| 59 | for (int i = 0; i < num_inputs; ++i) { |
| 60 | pattern.inputs[i] = {"*"}; |
| 61 | } |
| 62 | // Keep looking for nodes to remove until there are no more changes. |
| 63 | bool any_nodes_removed; |
| 64 | do { |
| 65 | any_nodes_removed = false; |
| 66 | std::map<string, string> inputs_to_rename; |
| 67 | GraphDef replaced_graph_def; |
| 68 | TF_RETURN_IF_ERROR(ReplaceMatchingOpTypes( |
| 69 | current_graph_def, pattern, |
| 70 | [&inputs_to_rename, &required_nodes, &any_nodes_removed]( |
| 71 | const NodeMatch& match, const std::set<string>& input_nodes, |
| 72 | const std::set<string>& output_nodes, |
| 73 | std::vector<NodeDef>* new_nodes) { |
| 74 | const NodeDef& replace_node = match.node; |
| 75 | // If this node is needed in the inputs or outputs don't replace |
| 76 | // it. |
| 77 | if (required_nodes.count(replace_node.name())) { |
| 78 | LOG(INFO) << "Skipping replacement for " << replace_node.name(); |
| 79 | CopyOriginalMatch(match, new_nodes); |
| 80 | return Status::OK(); |
| 81 | } |
| 82 | const NodeDef& input_node = match.inputs[0].node; |
| 83 | string target_name = input_node.name(); |
| 84 | for (const string& input : replace_node.input()) { |
| 85 | if (!input.compare(0, target_name.size(), target_name)) { |
| 86 | if (input.size() == target_name.size() || |
| 87 | input[target_name.size()] == ':') { |
| 88 | target_name = input; |