experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Compress constants in graph at the end of meta-optimizer.

Browse Source 
PiperOrigin-RevId: 235629151
This commit is contained in:
A. Unique TensorFlower 2019-02-25 17:02:36 -08:00 committed by TensorFlower Gardener
parent 84fdb93a93
commit 0956cbe9de
1 changed files with 16 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
tensorflow/core/grappler/optimizers/meta_optimizer.cc
View File

@ -17,6 +17,7 @@ limitations under the License.
#include "absl/strings/substitute.h" #include "absl/strings/substitute.h"
#include "tensorflow/core/common_runtime/function.h" #include "tensorflow/core/common_runtime/function.h"
#include "tensorflow/core/framework/function.pb.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/framework/versions.pb.h"
#include "tensorflow/core/graph/graph_constructor.h" #include "tensorflow/core/graph/graph_constructor.h"
#include "tensorflow/core/grappler/clusters/virtual_cluster.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 } // namespace
#define MK_OPT(NAME, VALUE) \ #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); 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. // Record graph optimization result.
optimization_results_.push_back(optimization_result); optimization_results_.push_back(optimization_result);