| 1640 | } |
| 1641 | |
| 1642 | void AddNode(const Node* node, |
| 1643 | const std::map<std::string, std::string> &prefix_map, |
| 1644 | GraphDef *gdef) { |
| 1645 | NodeDef *node_def = gdef->add_node(); |
| 1646 | *node_def = node->def(); |
| 1647 | std::string node_name = node_def->name(); |
| 1648 | auto it = prefix_map.find(node_name); |
| 1649 | if (it != prefix_map.end()) { |
| 1650 | node_name = it->second + "/" + node_name; |
| 1651 | *(node_def->mutable_name()) = node_name; |
| 1652 | } |
| 1653 | |
| 1654 | for (int i = 0; i < node_def->input_size(); i++) { |
| 1655 | std::string input_name = node_def->input(i); |
| 1656 | std::string input_node_name; |
| 1657 | std::string input_node_idx; |
| 1658 | size_t pos = input_name.find(":"); |
| 1659 | if (pos != string::npos) { |
| 1660 | input_node_name = input_name.substr(0, pos); |
| 1661 | input_node_idx = input_name.substr(pos); |
| 1662 | } else { |
| 1663 | input_node_name = input_name; |
| 1664 | } |
| 1665 | |
| 1666 | std::string control_prefix; |
| 1667 | if (input_node_name[0] == '^') { |
| 1668 | input_node_name = input_node_name.substr(1); |
| 1669 | control_prefix = "^"; |
| 1670 | } |
| 1671 | |
| 1672 | std::string prefix; |
| 1673 | auto it = prefix_map.find(input_node_name); |
| 1674 | if (it != prefix_map.end()) { |
| 1675 | prefix = it->second; |
| 1676 | } |
| 1677 | |
| 1678 | std::string new_input_name = control_prefix + prefix + |
| 1679 | "/" + input_node_name + input_node_idx; |
| 1680 | *node_def->mutable_input(i) = new_input_name; |
| 1681 | } |
| 1682 | } |
| 1683 | |
| 1684 | void PrintGraphForDebug(const SubGraph &worker_sub_graph, |
| 1685 | const std::vector<SubGraph> &ps_sub_graphs) |