| 189 | } |
| 190 | |
| 191 | std::unordered_set<int> SymbolicGradientBuilder::GetStopBackpropNodes( |
| 192 | const std::vector<bool>& reachable_nodes, |
| 193 | const std::unordered_set<int>& output_nodes) { |
| 194 | // Output nodes that get transitively consumed by other `outputs_` are stored |
| 195 | // in `internal_outputs`. |
| 196 | std::unordered_set<int> internal_outputs; |
| 197 | std::unordered_set<Node*> visited; |
| 198 | // Initialize `queue` for BFS traversal. Nodes in `queue` hold upcoming nodes |
| 199 | // along with the last Node in `output_` encountered along that path. If no |
| 200 | // `output_` node was encountered, pair.second will be nullptr. |
| 201 | std::deque<std::pair<Node*, Node*>> queue; |
| 202 | for (const Output& nout : inputs_) { |
| 203 | auto const& pair = visited.insert(nout.node()); |
| 204 | if (pair.second) { |
| 205 | queue.push_back(std::make_pair(nout.node(), static_cast<Node*>(nullptr))); |
| 206 | } |
| 207 | } |
| 208 | // BFS from nodes in 'inputs_' along out edges for the entire graph. Internal |
| 209 | // output nodes are recorded during the traversal. All nodes that are output |
| 210 | // nodes but not internal output nodes are considered the frontier of the |
| 211 | // output nodes, and thus our stop backprop nodes. |
| 212 | while (!queue.empty()) { |
| 213 | std::pair<Node*, Node*> p = queue.front(); |
| 214 | Node* n = p.first; |
| 215 | queue.pop_front(); |
| 216 | for (const Edge* e : n->out_edges()) { |
| 217 | // If a node is not reachable from outputs_, we can stop. |
| 218 | if (e->IsControlEdge() || !reachable_nodes[e->dst()->id()]) continue; |
| 219 | |
| 220 | auto const& pair = visited.insert(e->dst()); |
| 221 | if (pair.second) { |
| 222 | int node_id = e->dst()->id(); |
| 223 | Node* last_output_node = p.second; |
| 224 | if (output_nodes.find(node_id) != output_nodes.end()) { |
| 225 | // We reached an output node. |
| 226 | if (last_output_node != nullptr) { |
| 227 | // If we had already found an output node on this path so we mark |
| 228 | // it as an internal output. |
| 229 | internal_outputs.insert(last_output_node->id()); |
| 230 | } |
| 231 | // Mark this newly found output node to insert in the queue. |
| 232 | last_output_node = e->dst(); |
| 233 | } |
| 234 | queue.push_back(std::make_pair(e->dst(), last_output_node)); |
| 235 | } |
| 236 | } |
| 237 | } |
| 238 | // Finally, we set stop_backprop_nodes to all output_nodes that aren't also |
| 239 | // internal_outputs. |
| 240 | std::unordered_set<int> stop_backprop_nodes; |
| 241 | for (int output_node : output_nodes) { |
| 242 | if (internal_outputs.find(output_node) == internal_outputs.end()) { |
| 243 | stop_backprop_nodes.insert(output_node); |
| 244 | } |
| 245 | } |
| 246 | return stop_backprop_nodes; |
| 247 | } |
| 248 | |