| 558 | } |
| 559 | |
| 560 | Status SparsifyGather(const GraphDef& input_graph_def, |
| 561 | const TransformFuncContext& context, |
| 562 | GraphDef* output_graph_def) { |
| 563 | // clang-format off |
| 564 | const OpTypePattern gather_pattern = |
| 565 | {"Gather", |
| 566 | { |
| 567 | {"Identity", |
| 568 | { |
| 569 | {"Const|Variable|VariableV2"} |
| 570 | } |
| 571 | }, |
| 572 | {"*"}, |
| 573 | } |
| 574 | }; |
| 575 | const OpTypePattern gather_v2_pattern = |
| 576 | {"GatherV2", |
| 577 | { |
| 578 | {"Identity", |
| 579 | { |
| 580 | {"Const|Variable|VariableV2"} |
| 581 | } |
| 582 | }, |
| 583 | {"*"}, |
| 584 | // GatherV2's axis must be constant. |
| 585 | {"Const"}, |
| 586 | } |
| 587 | }; |
| 588 | // clang-format on |
| 589 | |
| 590 | GraphDef cleaned_input_graph_def; |
| 591 | RemoveAttributes(input_graph_def, {"_output_shapes"}, |
| 592 | &cleaned_input_graph_def); |
| 593 | |
| 594 | GraphDef temp_output; |
| 595 | |
| 596 | std::unique_ptr<BundleReader> ckpt_reader; |
| 597 | TF_RETURN_IF_ERROR(InitializeCheckpointReader(context, &ckpt_reader)); |
| 598 | |
| 599 | std::unique_ptr<std::unordered_map<string, string> > shapes_and_slices; |
| 600 | TF_RETURN_IF_ERROR( |
| 601 | ObtainVariableInfo(cleaned_input_graph_def, &shapes_and_slices)); |
| 602 | |
| 603 | TF_RETURN_IF_ERROR(SparsifyGatherInternal( |
| 604 | cleaned_input_graph_def, shapes_and_slices, context, gather_pattern, |
| 605 | ckpt_reader, &temp_output)); |
| 606 | |
| 607 | TF_RETURN_IF_ERROR(SparsifyGatherInternal(temp_output, shapes_and_slices, |
| 608 | context, gather_v2_pattern, |
| 609 | ckpt_reader, output_graph_def)); |
| 610 | |
| 611 | return Status::OK(); |
| 612 | } |
| 613 | |
| 614 | REGISTER_GRAPH_TRANSFORM("sparsify_gather", SparsifyGather); |
| 615 | |