| 424 | } |
| 425 | |
| 426 | Status PruneGraphDefInto(const tf2xla::Config& config, const GraphDef& in, |
| 427 | GraphDef* out) { |
| 428 | *out = in; |
| 429 | out->clear_node(); |
| 430 | |
| 431 | // Tensors needed for feeding. |
| 432 | std::set<std::pair<string, int>> feed_tensors; |
| 433 | for (const tf2xla::Feed& feed : config.feed()) { |
| 434 | feed_tensors.insert( |
| 435 | std::make_pair(feed.id().node_name(), feed.id().output_index())); |
| 436 | } |
| 437 | |
| 438 | // Maps node name to reachability. |
| 439 | std::unordered_map<string, std::pair<bool, const NodeDef*>> node_by_name; |
| 440 | for (const NodeDef& node : in.node()) { |
| 441 | node_by_name[node.name()] = std::pair<bool, const NodeDef*>(false, &node); |
| 442 | } |
| 443 | |
| 444 | // Traverse. |
| 445 | std::queue<string> name_queue; |
| 446 | for (int i = 0; i < config.fetch_size(); ++i) { |
| 447 | name_queue.push(config.fetch(i).id().node_name()); |
| 448 | } |
| 449 | while (!name_queue.empty()) { |
| 450 | const string name = name_queue.front(); |
| 451 | name_queue.pop(); |
| 452 | |
| 453 | auto find_it = node_by_name.find(name); |
| 454 | if (find_it == node_by_name.end()) { |
| 455 | return errors::InvalidArgument("While pruning graph, node ", name, |
| 456 | " needed but not found in the graph."); |
| 457 | } |
| 458 | auto& map_entry = find_it->second; |
| 459 | if (map_entry.first) { |
| 460 | continue; |
| 461 | } |
| 462 | map_entry.first = true; |
| 463 | |
| 464 | // Push input nodes of the currently visited node to name_queue. |
| 465 | for (const string& in_edge : map_entry.second->input()) { |
| 466 | auto id = ParseTensorName(in_edge); |
| 467 | const string node_name = string(id.first); |
| 468 | if (feed_tensors.find(std::make_pair(node_name, id.second)) == |
| 469 | feed_tensors.end()) { |
| 470 | name_queue.push(node_name); |
| 471 | } else { |
| 472 | // The input tensor is from an edge that is being fed. Therefore, |
| 473 | // we skip recursing down that edge, to avoid requiring nodes that |
| 474 | // may not be needed (note that the input node may still be added |
| 475 | // to name_queue later if one of its output edges is not being fed). |
| 476 | } |
| 477 | } |
| 478 | } |
| 479 | |
| 480 | // Copy over, preserving order of original and only nodes that are reachable |
| 481 | // from the fetches. |
| 482 | out->mutable_node()->Reserve(in.node_size()); |
| 483 | for (const NodeDef& node : in.node()) { |