diff options
| author |  gracehoney <31743510+aaroey@users.noreply.github.com> | 2018-04-27 13:07:16 -0700 |
|---|---|---|
| committer |  Benoit Steiner <benoitsteiner@users.noreply.github.com> | 2018-04-27 13:07:16 -0700 |
| commit | 3e582feec7b7e1b71d5a4b590edc1e4d4e4a3126 (patch) | |
| tree | f8fecc5a818142316c525a469a0827fe309d9e00 | |
| parent | 32d0063f0a3f7adbc4c0f6ef59b96504a7208052 (diff) | |
Roll forward the custom optimizers change (#18742)
| -rw-r--r-- | tensorflow/core/grappler/optimizers/meta_optimizer.cc | 17 |
1 files changed, 15 insertions, 2 deletions
diff --git a/tensorflow/core/grappler/optimizers/meta_optimizer.cc b/tensorflow/core/grappler/optimizers/meta_optimizer.cc index 2edc4da9dc..5230177dca 100644 --- a/tensorflow/core/grappler/optimizers/meta_optimizer.cc +++ b/tensorflow/core/grappler/optimizers/meta_optimizer.cc @@ -160,13 +160,26 @@ Status MetaOptimizer::InitializeOptimizersByName( VLOG(2) << "Can't register an optimizer by name: " << optimizer_name; } } + for (const auto& optimizer_config : cfg_.custom_optimizers()) { + auto custom_optimizer = CustomGraphOptimizerRegistry::CreateByNameOrNull( + optimizer_config.name()); + if (custom_optimizer) { + VLOG(2) << "Registered custom configurable graph optimizer: " + << optimizer_config.name(); + TF_RETURN_IF_ERROR(custom_optimizer->Init(&optimizer_config)); + optimizers->push_back(std::move(custom_optimizer)); + } else { + VLOG(2) << "Can't register an optimizer by name: " + << optimizer_config.name(); + } + } return Status::OK(); } Status MetaOptimizer::OptimizeGraph(Cluster* cluster, const GrapplerItem& item, GraphDef* optimized_graph) { std::vector<std::unique_ptr<GraphOptimizer>> optimizers; - if (cfg_.optimizers().empty()) { + if (cfg_.optimizers().empty() && cfg_.custom_optimizers().empty()) { TF_RETURN_IF_ERROR(InitializeOptimizers(&optimizers)); } else { TF_RETURN_IF_ERROR(InitializeOptimizersByName(&optimizers)); @@ -337,7 +350,7 @@ bool MetaOptimizerEnabled(const RewriterConfig& cfg) { cfg.auto_parallel().enable() || cfg.memory_optimization() != RewriterConfig::NO_MEM_OPT || cfg.debug_stripper() == RewriterConfig::ON || - !cfg.optimizers().empty(); + !cfg.optimizers().empty() || !cfg.custom_optimizers().empty(); } Status RunMetaOptimizer(const GrapplerItem& item, const RewriterConfig& cfg, |