| 668 | } |
| 669 | |
| 670 | std::shared_ptr<Node> Model::AddNode(Node::Factory factory, const string& name, |
| 671 | const string& output_name) { |
| 672 | // The name captures the sequence of iterators joined by `::`. We use the full |
| 673 | // sequence as the key in the lookup table, but only the last element of the |
| 674 | // sequence as the name node. |
| 675 | std::vector<string> tokens = |
| 676 | str_util::Split(name, ':', str_util::SkipEmpty()); |
| 677 | // The output name might contain an index. We need to strip it to make it |
| 678 | // possible for the model to successfully identify the output node. |
| 679 | string sanitized_output_name = output_name; |
| 680 | if (str_util::EndsWith(output_name, "]")) { |
| 681 | sanitized_output_name = output_name.substr(0, output_name.rfind('[')); |
| 682 | } |
| 683 | std::shared_ptr<Node> output; |
| 684 | mutex_lock l(mu_); |
| 685 | auto it = lookup_table_.find(sanitized_output_name); |
| 686 | if (it != lookup_table_.end()) { |
| 687 | output = it->second; |
| 688 | } |
| 689 | std::shared_ptr<Node> node = factory({id_counter_++, tokens.back(), output}); |
| 690 | if (!output_) { |
| 691 | output_ = node; |
| 692 | } |
| 693 | if (output) { |
| 694 | VLOG(3) << "Adding " << node->long_name() << " as input for " |
| 695 | << output->long_name(); |
| 696 | output->add_input(node); |
| 697 | } else { |
| 698 | VLOG(3) << "Adding " << node->long_name(); |
| 699 | } |
| 700 | collect_resource_usage_ = |
| 701 | collect_resource_usage_ || node->has_tunable_parameters(); |
| 702 | lookup_table_.insert(std::make_pair(name, node)); |
| 703 | return node; |
| 704 | } |
| 705 | |
| 706 | void Model::AddProcessingTime(const string& name, int64 delta) { |
| 707 | tf_shared_lock l(mu_); |