From 0956cbe9dedebc168ac3a52b6fdaee826a604a15 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 25 Feb 2019 17:02:36 -0800
Subject: [PATCH] Compress constants in graph at the end of meta-optimizer.

PiperOrigin-RevId: 235629151
---
 .../core/grappler/optimizers/meta_optimizer.cc   | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/tensorflow/core/grappler/optimizers/meta_optimizer.cc b/tensorflow/core/grappler/optimizers/meta_optimizer.cc
index 36d68a7b0fb..fd420e9f7fc 100644
--- a/tensorflow/core/grappler/optimizers/meta_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/meta_optimizer.cc
@@ -17,6 +17,7 @@ limitations under the License.
 #include "absl/strings/substitute.h"
 #include "tensorflow/core/common_runtime/function.h"
 #include "tensorflow/core/framework/function.pb.h"
+#include "tensorflow/core/framework/tensor_util.h"
 #include "tensorflow/core/framework/versions.pb.h"
 #include "tensorflow/core/graph/graph_constructor.h"
 #include "tensorflow/core/grappler/clusters/virtual_cluster.h"
@@ -102,6 +103,18 @@ uint64 DeadlineMicroSeconds(const RewriterConfig& cfg) {
   }
 }
 
+Status CompressConstants(GraphDef* graph) {
+  for (int i = 0; i < graph->node_size(); ++i) {
+    NodeDef* node = graph->mutable_node(i);
+    if ((IsConstant(*node) || IsHostConstant(*node)) &&
+        HasNodeAttr(*node, "value")) {
+      AttrValue& attr_val = (*node->mutable_attr())["value"];
+      tensor::CompressTensorProtoInPlace(attr_val.mutable_tensor());
+    }
+  }
+  return Status::OK();
+}
+
 }  // namespace
 
 #define MK_OPT(NAME, VALUE) \
@@ -432,6 +445,9 @@ Status MetaOptimizer::OptimizeGraph(Cluster* cluster, const GrapplerItem& item,
     RUN_OPTIMIZER_OR_RETURN_IF_ERROR(sa_optimizer);
   }
 
+  // Compress the constants in the final graph.
+  TF_RETURN_IF_ERROR(CompressConstants(optimized_graph));
+
   // Record graph optimization result.
   optimization_results_.push_back(optimization_result);