| 382 | } |
| 383 | |
| 384 | Status GetInitOp(const string& export_dir, const MetaGraphDef& meta_graph_def, |
| 385 | string* init_op_name) { |
| 386 | const auto& sig_def_map = meta_graph_def.signature_def(); |
| 387 | const auto& init_op_sig_it = |
| 388 | meta_graph_def.signature_def().find(kSavedModelInitOpSignatureKey); |
| 389 | if (init_op_sig_it != sig_def_map.end()) { |
| 390 | const auto& sig_def_outputs = init_op_sig_it->second.outputs(); |
| 391 | const auto& sig_def_outputs_it = |
| 392 | sig_def_outputs.find(kSavedModelInitOpSignatureKey); |
| 393 | if (sig_def_outputs_it == sig_def_outputs.end()) { |
| 394 | return errors::FailedPrecondition("Could not find output ", |
| 395 | kSavedModelInitOpSignatureKey); |
| 396 | } |
| 397 | *init_op_name = sig_def_outputs_it->second.name(); |
| 398 | return Status::OK(); |
| 399 | } |
| 400 | |
| 401 | const auto& collection_def_map = meta_graph_def.collection_def(); |
| 402 | string init_op_collection_key; |
| 403 | if (collection_def_map.find(kSavedModelMainOpKey) != |
| 404 | collection_def_map.end()) { |
| 405 | init_op_collection_key = kSavedModelMainOpKey; |
| 406 | } else { |
| 407 | init_op_collection_key = kSavedModelLegacyInitOpKey; |
| 408 | } |
| 409 | |
| 410 | const auto init_op_it = collection_def_map.find(init_op_collection_key); |
| 411 | if (init_op_it != collection_def_map.end()) { |
| 412 | if (init_op_it->second.node_list().value_size() != 1) { |
| 413 | return errors::FailedPrecondition( |
| 414 | strings::StrCat("Expected exactly one main op in : ", export_dir)); |
| 415 | } |
| 416 | *init_op_name = init_op_it->second.node_list().value(0); |
| 417 | } |
| 418 | return Status::OK(); |
| 419 | } |
| 420 | |
| 421 | namespace { |
| 422 | Status ValidateNode(const NodeDef& node) { |
no test coverage detected