| 475 | } |
| 476 | |
| 477 | Status RenameNodeInputs(const GraphDef& input_graph_def, |
| 478 | const std::map<string, string>& inputs_to_rename, |
| 479 | const std::unordered_set<string>& nodes_to_ignore, |
| 480 | GraphDef* output_graph_def) { |
| 481 | std::map<string, std::vector<std::pair<string, string>>> |
| 482 | canonical_inputs_to_rename; |
| 483 | for (const auto& input_to_rename : inputs_to_rename) { |
| 484 | canonical_inputs_to_rename[NodeNameFromInput(input_to_rename.first)] |
| 485 | .push_back({input_to_rename.first, input_to_rename.second}); |
| 486 | } |
| 487 | |
| 488 | output_graph_def->Clear(); |
| 489 | for (const NodeDef& node : input_graph_def.node()) { |
| 490 | NodeDef* new_node = output_graph_def->mutable_node()->Add(); |
| 491 | *new_node = node; |
| 492 | new_node->mutable_input()->Clear(); |
| 493 | for (const string& input_name : node.input()) { |
| 494 | std::set<string> already_visited; |
| 495 | string new_input_name = input_name; |
| 496 | while ( |
| 497 | canonical_inputs_to_rename.count(NodeNameFromInput(new_input_name))) { |
| 498 | string input_node_name = NodeNameFromInput(new_input_name); |
| 499 | if (already_visited.count(input_node_name)) { |
| 500 | return errors::InvalidArgument( |
| 501 | "RenameNodeInputs argument contains a cycle for ", |
| 502 | input_node_name); |
| 503 | } |
| 504 | already_visited.insert(input_node_name); |
| 505 | if (nodes_to_ignore.count(node.name())) { |
| 506 | break; |
| 507 | } |
| 508 | bool any_match_found = false; |
| 509 | for (const std::pair<string, string>& input_to_rename : |
| 510 | canonical_inputs_to_rename.at(input_node_name)) { |
| 511 | const string& source_name = input_to_rename.first; |
| 512 | const string& dest_name = input_to_rename.second; |
| 513 | bool is_match; |
| 514 | string match_name; |
| 515 | if (str_util::EndsWith(source_name, ":*")) { |
| 516 | is_match = true; |
| 517 | string prefix; |
| 518 | string unused_node_name; |
| 519 | string suffix; |
| 520 | NodeNamePartsFromInput(new_input_name, &prefix, &unused_node_name, |
| 521 | &suffix); |
| 522 | match_name = prefix + dest_name + suffix; |
| 523 | } else { |
| 524 | is_match = (CanonicalInputName(source_name) == |
| 525 | CanonicalInputName(new_input_name)); |
| 526 | match_name = dest_name; |
| 527 | } |
| 528 | if (is_match) { |
| 529 | new_input_name = match_name; |
| 530 | any_match_found = true; |
| 531 | } |
| 532 | } |
| 533 | if (!any_match_found) { |
| 534 | break; |