Gets a map from variable name to variable value.
| 120 | |
| 121 | // Gets a map from variable name to variable value. |
| 122 | Status GetVariableNameToTensorMap( |
| 123 | Session* session, |
| 124 | const std::unordered_map<string, NodeDef*>& name_to_node_map, |
| 125 | std::unordered_set<string> variable_names_set, |
| 126 | std::unordered_map<string, Tensor>* variable_name_to_value_map) { |
| 127 | if (variable_names_set.empty()) { |
| 128 | return Status::OK(); |
| 129 | } |
| 130 | std::vector<string> variable_names; |
| 131 | variable_names.reserve(variable_names_set.size()); |
| 132 | std::vector<string> tensor_names; |
| 133 | tensor_names.reserve(variable_names_set.size()); |
| 134 | for (const string& node_name : variable_names_set) { |
| 135 | variable_names.push_back(node_name); |
| 136 | NodeDef* node_def = name_to_node_map.at(node_name); |
| 137 | if (node_def->op() == "VarHandleOp") { |
| 138 | // If this is a resource variable, we have to run the corresponding |
| 139 | // ReadVariableOp. |
| 140 | tensor_names.push_back(node_name + "/Read/ReadVariableOp:0"); |
| 141 | } else { |
| 142 | tensor_names.push_back(node_name + ":0"); |
| 143 | } |
| 144 | } |
| 145 | std::vector<Tensor> outputs; |
| 146 | TF_RETURN_IF_ERROR( |
| 147 | session->Run(/* inputs */ {}, tensor_names, /* targets */ {}, &outputs)); |
| 148 | for (size_t i = 0; i < variable_names.size(); i++) { |
| 149 | (*variable_name_to_value_map)[variable_names[i]] = outputs[i]; |
| 150 | } |
| 151 | return Status::OK(); |
| 152 | } |
| 153 | |
| 154 | // Converts a Variable NodeDef into a Constant NodeDef. |
| 155 | void ConvertVariableToConstant(const NodeDef& variable_node, |