TODO(zhifengc): Maybe this should be the default Graph::AsGraphDef. and stash the original NodeDef name as an attr for documentation purpose.
| 2241 | // and stash the original NodeDef name as an attr for documentation |
| 2242 | // purpose. |
| 2243 | void ToGraphDef(const Graph* g, GraphDef* gdef, bool pretty) { |
| 2244 | // We visit nodes in forward topological sort order, which is a |
| 2245 | // possible execution order of the graph. |
| 2246 | gtl::InlinedVector<const Edge*, 4> inputs; |
| 2247 | gdef->Clear(); |
| 2248 | *gdef->mutable_versions() = g->versions(); |
| 2249 | |
| 2250 | std::vector<Node*> start_nodes; |
| 2251 | for (Node* n : g->nodes()) { |
| 2252 | if (n->out_edges().empty()) { |
| 2253 | start_nodes.push_back(n); |
| 2254 | } |
| 2255 | } |
| 2256 | |
| 2257 | ReverseDFSFrom(*g, start_nodes, nullptr, [gdef, pretty, &inputs](Node* n) { |
| 2258 | if (!n->IsOp()) return; |
| 2259 | NodeDef* ndef = gdef->add_node(); |
| 2260 | ndef->set_name(NewName(n, pretty)); |
| 2261 | ndef->set_op(n->type_string()); |
| 2262 | for (const auto& attr : n->attrs()) { |
| 2263 | (*ndef->mutable_attr())[attr.first] = attr.second; |
| 2264 | } |
| 2265 | |
| 2266 | if (!n->assigned_device_name().empty()) { |
| 2267 | ndef->set_device(n->assigned_device_name()); |
| 2268 | } else { |
| 2269 | ndef->set_device(n->requested_device()); |
| 2270 | } |
| 2271 | |
| 2272 | inputs.clear(); |
| 2273 | inputs.resize(n->num_inputs()); |
| 2274 | for (const Edge* e : n->in_edges()) { |
| 2275 | if (e->IsControlEdge()) { |
| 2276 | inputs.push_back(e); |
| 2277 | } else { |
| 2278 | if (inputs[e->dst_input()] == nullptr) { |
| 2279 | inputs[e->dst_input()] = e; |
| 2280 | } else { |
| 2281 | LOG(WARNING) << "Malformed graph node. multiple input edges: " |
| 2282 | << n->DebugString(); |
| 2283 | } |
| 2284 | } |
| 2285 | } |
| 2286 | // node->name() is merely NodeDef::name, which are not guaranteed |
| 2287 | // to be unique and stable after optimization rewrites. Therefore, |
| 2288 | // we use "n<node id>" instead. |
| 2289 | for (const Edge* e : inputs) { |
| 2290 | if (e == nullptr) { |
| 2291 | ndef->add_input("unknown"); |
| 2292 | continue; |
| 2293 | } |
| 2294 | const string srcname = NewName(e->src(), pretty); |
| 2295 | if (!e->src()->IsOp()) { |
| 2296 | } else if (e->IsControlEdge()) { |
| 2297 | ndef->add_input(strings::StrCat("^", srcname)); |
| 2298 | } else if (e->src_output() == 0) { |
| 2299 | ndef->add_input(srcname); |
| 2300 | } else { |