Fix bugs in model pruner.

PiperOrigin-RevId: 195298816

diff --git a/tensorflow/core/grappler/optimizers/model_pruner.cc b/tensorflow/core/grappler/optimizers/model_pruner.cc
index 3311e97..36eab49 100644 (file)
--- a/tensorflow/core/grappler/optimizers/model_pruner.cc
+++ b/tensorflow/core/grappler/optimizers/model_pruner.cc
@@ -70,6 +70,7 @@ Status ModelPruner::Optimize(Cluster* cluster, const GrapplerItem& item,
     }
     // Try to keep the nodes ordered somewhat topologically since this helps
     // further optimizations perform better.
+    runnable_item.graph.mutable_node()->Reserve(keep.size());
     for (int i = keep.size() - 1; i >= 0; --i) {
       *runnable_item.graph.add_node() = *keep[i];
     }
@@ -113,6 +114,7 @@ Status ModelPruner::Optimize(Cluster* cluster, const GrapplerItem& item,
     }
   }
 
+  pruned_graph->Clear();
   *pruned_graph->mutable_library() = item.graph.library();
   *pruned_graph->mutable_versions() = item.graph.versions();
 
@@ -122,6 +124,7 @@ Status ModelPruner::Optimize(Cluster* cluster, const GrapplerItem& item,
   }
 
   const bool fetches_are_known = !item.fetch.empty();
+  pruned_graph->mutable_node()->Reserve(runnable_item.graph.node_size());
   for (auto& node : runnable_item.graph.node()) {
     if (!fetches_are_known ||
         nodes_to_delete.find(&node) == nodes_to_delete.end()) {
@@ -134,6 +137,7 @@ Status ModelPruner::Optimize(Cluster* cluster, const GrapplerItem& item,
   VLOG(1) << "Pruned " << nodes_to_delete.size()
           << " nodes from the graph. The graph now contains "
           << pruned_graph->node_size() << " nodes.";
+  CHECK_LE(pruned_graph->node_size(), item.graph.node_size());
 
   return Status::OK();
 }