| 2115 | } |
| 2116 | |
| 2117 | Status TrainGraphPartitioner::SplitGraph( |
| 2118 | SubGraph *worker_sub_graph, |
| 2119 | std::vector<SubGraph> *ps_sub_graphs, |
| 2120 | bool merge_ps_graph) |
| 2121 | { |
| 2122 | |
| 2123 | RETURN_IF_NOT_OK(SplitGraphInternal(ps_sub_graphs, worker_sub_graph, true)); |
| 2124 | |
| 2125 | for (auto &sub_graph : *ps_sub_graphs) { |
| 2126 | sub_graph.CompleteVariables(opts_); |
| 2127 | } |
| 2128 | |
| 2129 | if (merge_ps_graph) { |
| 2130 | std::vector<SubGraph> merged_ps_sub_graphs; |
| 2131 | auto status = MergePsGraphs(*worker_sub_graph, |
| 2132 | *ps_sub_graphs, |
| 2133 | &merged_ps_sub_graphs); |
| 2134 | if (status.ok()) { |
| 2135 | ps_sub_graphs->swap(merged_ps_sub_graphs); |
| 2136 | } else { |
| 2137 | LOG(FATAL) << "Merge ps graph error, " << status.error_message(); |
| 2138 | } |
| 2139 | } |
| 2140 | |
| 2141 | return Status::OK(); |
| 2142 | } |
| 2143 | |
| 2144 | Status InferGraphPartitioner::SplitGraph( |
| 2145 | SubGraph *main_sub_graph, |