| 1700 | } |
| 1701 | |
| 1702 | Status SavedModelOptimizer::AddVariableInitSubGraph() { |
| 1703 | // Add a init_op to initialize storage like redis. |
| 1704 | NodeDef init_op_def; |
| 1705 | init_op_def.set_name(GetInitNodeName()); |
| 1706 | init_op_def.set_op("KvInit"); |
| 1707 | |
| 1708 | std::vector<std::string> feature_names; |
| 1709 | for (Node* node : graph_.nodes()) { |
| 1710 | // Add resource name as feature_names attr |
| 1711 | if (node->op_def().name() == "KvResourceImportV2") { |
| 1712 | feature_names.push_back(node->name()); |
| 1713 | } |
| 1714 | } |
| 1715 | AddNodeAttr("feature_names", feature_names, &init_op_def); |
| 1716 | |
| 1717 | Status status; |
| 1718 | Node* init_op_node = graph_.AddNode(init_op_def, &status); |
| 1719 | return status; |
| 1720 | } |
| 1721 | |
| 1722 | namespace { |
| 1723 | void GetResourceNodeFromRestoreOpInputs( |