| 191 | } |
| 192 | |
| 193 | void TFStats::AddGraph(std::unique_ptr<GraphDef> graph) { |
| 194 | std::map<string, const NodeDef*> node_defs; |
| 195 | bool node_added = false; |
| 196 | for (const NodeDef& node : graph->node()) { |
| 197 | if (nodes_map_.find(node.name()) != nodes_map_.end()) { |
| 198 | continue; |
| 199 | } |
| 200 | node_added = true; |
| 201 | size_t num_nodes = nodes_map_.size(); |
| 202 | nodes_map_[node.name()] = std::unique_ptr<TFGraphNode>( |
| 203 | new TFGraphNode(&node, num_nodes, &nodes_map_)); |
| 204 | node_defs[node.name()] = &node; |
| 205 | } |
| 206 | for (auto it = node_defs.begin(); it != node_defs.end(); it++) { |
| 207 | TFGraphNode* node = nodes_map_.at(it->first).get(); |
| 208 | for (int i = 0; i < it->second->input_size(); ++i) { |
| 209 | string node_input = it->second->input(i); |
| 210 | int output_idx = 0; |
| 211 | // input name format can be: "^node:src_output" |
| 212 | // if not :src_output, then it's the first one (further verify?) |
| 213 | auto prefix_pos = node_input.find(":"); |
| 214 | if (prefix_pos != node_input.npos) { |
| 215 | std::vector<string> input_parts = str_util::Split(node_input, ":"); |
| 216 | CHECK(input_parts.size() == 2) |
| 217 | << "Unknown NodeDef.input format: " << node_input; |
| 218 | node_input = input_parts[0]; |
| 219 | CHECK(strings::safe_strto32(input_parts[1], &output_idx)) |
| 220 | << "Failed to parse integer: " << output_idx; |
| 221 | } |
| 222 | if (node_input.substr(0, 1) == "^") { |
| 223 | node_input = node_input.substr(1); |
| 224 | } |
| 225 | // Delay input TFGraphNode retrieval as late as possible. |
| 226 | // In long run, when we have TensorFlow runtime graph, the |
| 227 | // graph connection should be dynamic and per-step. |
| 228 | node->AddInput(node_input, output_idx, i); |
| 229 | } |
| 230 | } |
| 231 | if (node_added) { |
| 232 | graph_view_.reset(nullptr); |
| 233 | scope_view_.reset(nullptr); |
| 234 | op_view_.reset(nullptr); |
| 235 | code_view_.reset(nullptr); |
| 236 | } |
| 237 | } |
| 238 | |
| 239 | void TFStats::AddOpLogProto(std::unique_ptr<OpLogProto> op_log) { |
| 240 | if (!op_log) { |