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

For custom optimizers, keep a copy of Session ConfigProto_.

Browse Source 
PiperOrigin-RevId: 292974333
Change-Id: Ia49fae2149aea5e3e4fec5ba9d4b995406a6a1e4
This commit is contained in:
A. Unique TensorFlower 2020-02-03 12:08:50 -08:00 committed by TensorFlower Gardener
parent ab94de29f5
commit be453c9869
2 changed files with 14 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
tensorflow/core/grappler/optimizers/custom_graph_optimizer.h
View File

@ -18,6 +18,7 @@ limitations under the License.
#include "tensorflow/core/grappler/optimizers/graph_optimizer.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/protobuf/config.pb.h"
#include "tensorflow/core/protobuf/rewriter_config.pb.h"
namespace tensorflow {
@ -29,6 +30,15 @@ class CustomGraphOptimizer : public GraphOptimizer {
virtual ~CustomGraphOptimizer() {}
virtual Status Init(const tensorflow::RewriterConfig_CustomGraphOptimizer*
config = nullptr) = 0;
// Populates ConfigProto on which the Session is run prior to running Init.
Status InitWithConfig(
const ConfigProto& config_proto,
const tensorflow::RewriterConfig_CustomGraphOptimizer* config = nullptr) {
config_proto_ = config_proto;
return this->Init(config);
}
ConfigProto config_proto_;
};
} // end namespace grappler

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

@ -299,8 +299,8 @@ Status MetaOptimizer::InitializeOptimizersByName(
if (custom_optimizer) {
VLOG(2) << "Registered custom graph optimizer: " << optimizer_name;
TF_RETURN_IF_ERROR(custom_optimizer->Init(
GetCustomGraphOptimizerConfig(optimizer_name)));
TF_RETURN_IF_ERROR(custom_optimizer->InitWithConfig(
config_proto_, GetCustomGraphOptimizerConfig(optimizer_name)));
optimizers->push_back(std::move(custom_optimizer));
initialized_custom_optimizers.insert(optimizer_name);
} else {
@ -326,7 +326,8 @@ Status MetaOptimizer::InitializeCustomGraphOptimizers(
if (custom_optimizer) {
VLOG(2) << "Registered custom configurable graph optimizer: "
<< optimizer_config.name();
TF_RETURN_IF_ERROR(custom_optimizer->Init(&optimizer_config));
TF_RETURN_IF_ERROR(
custom_optimizer->InitWithConfig(config_proto_, &optimizer_config));
optimizers->push_back(std::move(custom_optimizer));
} else {
// If there are no custom optimizers with given name, try to initialize a