diff options
4 files changed, 10 insertions, 7 deletions
diff --git a/tensorflow/core/grappler/optimizers/meta_optimizer.cc b/tensorflow/core/grappler/optimizers/meta_optimizer.cc index b1f31ad0d0..c55f479451 100644 --- a/tensorflow/core/grappler/optimizers/meta_optimizer.cc +++ b/tensorflow/core/grappler/optimizers/meta_optimizer.cc @@ -91,7 +91,8 @@ std::unique_ptr<GraphOptimizer> MetaOptimizer::MakeNewOptimizer( MK_OPT("dependency", new DependencyOptimizer(cfg_.dependency_optimization())); MK_OPT("debug_stripper", new DebugStripper()); MK_OPT("scoped_allocator", - new ScopedAllocatorOptimizer(cfg_.scoped_allocator_opts())); + new ScopedAllocatorOptimizer(cfg_.scoped_allocator_optimization(), + cfg_.scoped_allocator_opts())); return std::unique_ptr<GraphOptimizer>(); } @@ -150,8 +151,8 @@ Status MetaOptimizer::InitializeOptimizers( new AutoParallel(cfg_.auto_parallel().num_replicas())); } if (cfg_.scoped_allocator_optimization()) { - optimizers->emplace_back( - new ScopedAllocatorOptimizer(cfg_.scoped_allocator_opts())); + optimizers->emplace_back(new ScopedAllocatorOptimizer( + cfg_.scoped_allocator_optimization(), cfg_.scoped_allocator_opts())); } return Status::OK(); } diff --git a/tensorflow/core/grappler/optimizers/scoped_allocator_optimizer.cc b/tensorflow/core/grappler/optimizers/scoped_allocator_optimizer.cc index cceef4098d..275568e464 100644 --- a/tensorflow/core/grappler/optimizers/scoped_allocator_optimizer.cc +++ b/tensorflow/core/grappler/optimizers/scoped_allocator_optimizer.cc @@ -650,7 +650,8 @@ class UnaryElementwiseRewriter : public ScopedAllocatorOptimizer::Rewriter { }; ScopedAllocatorOptimizer::ScopedAllocatorOptimizer( - const ScopedAllocatorOptions& opts) { + RewriterConfig::Toggle opt_level, const ScopedAllocatorOptions& opts) + : opt_level_(opt_level) { VLOG(1) << "ScopedAllocatorOptimizer::ScopedAllocatorOptimizer"; Rewriter* r = new UnaryElementwiseRewriter(); to_delete_.push_back(r); diff --git a/tensorflow/core/grappler/optimizers/scoped_allocator_optimizer.h b/tensorflow/core/grappler/optimizers/scoped_allocator_optimizer.h index ab4d444595..13589f536c 100644 --- a/tensorflow/core/grappler/optimizers/scoped_allocator_optimizer.h +++ b/tensorflow/core/grappler/optimizers/scoped_allocator_optimizer.h @@ -32,7 +32,8 @@ class ScopedAllocatorOptimizer; // movement and consolidate some kinds of Ops. class ScopedAllocatorOptimizer : public GraphOptimizer { public: - explicit ScopedAllocatorOptimizer(const ScopedAllocatorOptions& opts); + ScopedAllocatorOptimizer(RewriterConfig::Toggle opt_level, + const ScopedAllocatorOptions& opts); ~ScopedAllocatorOptimizer() override; string name() const override { return "scoped_allocator_optimizer"; } diff --git a/tensorflow/core/grappler/optimizers/scoped_allocator_optimizer_test.cc b/tensorflow/core/grappler/optimizers/scoped_allocator_optimizer_test.cc index 3a2859dc5f..89847f83d4 100644 --- a/tensorflow/core/grappler/optimizers/scoped_allocator_optimizer_test.cc +++ b/tensorflow/core/grappler/optimizers/scoped_allocator_optimizer_test.cc @@ -115,7 +115,7 @@ TEST_F(ScopedAllocatorOptimizerTest, UnaryRewriteOnly) { ScopedAllocatorOptions opts; opts.add_enable_op("Abs"); - ScopedAllocatorOptimizer sao(opts); + ScopedAllocatorOptimizer sao(RewriterConfig::ON, opts); ScopedAllocatorOptimizer::OpNameSet ons; ons.insert("Abs"); @@ -199,7 +199,7 @@ TEST_F(ScopedAllocatorOptimizerTest, UnaryExecute) { // b + c == -4, -4, 3, 2 for (int oi = 0; oi < outputs.size(); ++oi) { for (int i = 0; i < outputs[oi].NumElements(); ++i) { - VLOG(0) << "output vec " << oi << " index " << i << " = " + VLOG(1) << "output vec " << oi << " index " << i << " = " << outputs[oi].flat<float>()(i); } if (oi == 0) { |