| 50 | }; |
| 51 | |
| 52 | class SubGraph { |
| 53 | public: |
| 54 | typedef std::vector<std::pair<std::string, const Edge*>> KeyEdgeVec; |
| 55 | SubGraph(const std::string &loc) : |
| 56 | loc_(loc), only_variable_(false), use_send_recv_(false) {} |
| 57 | SubGraph() : only_variable_(false), use_send_recv_(false) {} |
| 58 | const std::string& GetLoc() const { return loc_; } |
| 59 | void SetLoc(const std::string &loc) { loc_ = loc; } |
| 60 | |
| 61 | void AddNode(const Node* node) { |
| 62 | if (node_id_set_.find(node->id()) != node_id_set_.end()) { |
| 63 | return; |
| 64 | } |
| 65 | nodes_.push_back(node); |
| 66 | node_id_set_.insert(node->id()); |
| 67 | } |
| 68 | |
| 69 | void AddInputEdge(const std::string& feed_name, const Edge* edge) { |
| 70 | input_edges_.push_back({feed_name, edge}); |
| 71 | } |
| 72 | |
| 73 | void AddInputs(const std::vector<std::pair<std::string, const Edge*> > &input_edges) { |
| 74 | input_edges_.insert(input_edges_.end(), input_edges.begin(), input_edges.end()); |
| 75 | } |
| 76 | |
| 77 | void AddOutputEdge(const std::string& fetch_name, const Edge* edge) { |
| 78 | output_edges_.push_back({fetch_name, edge}); |
| 79 | } |
| 80 | void AddOutputs(const std::vector<std::pair<std::string, const Edge*> > &output_edges) { |
| 81 | output_edges_.insert(output_edges_.end(), output_edges.begin(), output_edges.end()); |
| 82 | } |
| 83 | |
| 84 | bool NodeExisted(int id) { |
| 85 | return node_id_set_.find(id) != node_id_set_.end(); |
| 86 | } |
| 87 | void SetOnlyVariable() { only_variable_ = true;} |
| 88 | bool IsOnlyVariable() const { return only_variable_;} |
| 89 | const KeyEdgeVec& GetInputEdges() const { return input_edges_; } |
| 90 | const KeyEdgeVec& GetOutputEdges() const { return output_edges_; } |
| 91 | |
| 92 | const std::vector<const Node*>& GetNodes() const { return nodes_; } |
| 93 | |
| 94 | GraphDef& GetGraphDef() { return graph_def_; } |
| 95 | |
| 96 | const GraphDef& GetGraphDef() const { return graph_def_; } |
| 97 | |
| 98 | void Extend(const SubGraph &sub_graph) { |
| 99 | for (const Node* node : sub_graph.GetNodes()) { |
| 100 | if (loc_ == "") { |
| 101 | loc_ = sub_graph.GetLoc(); |
| 102 | } |
| 103 | AddNode(node); |
| 104 | } |
| 105 | } |
| 106 | |
| 107 | void SetGraphHandle(const std::string &handle) { |
| 108 | graph_handle_ = handle; |
| 109 | } |