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

[SavedModel C++] Avoid copying the MetaGraphDef in FindMetaGraphDef().

Browse Source 
PiperOrigin-RevId: 265502985
This commit is contained in:
Derek Murray 2019-08-26 11:35:46 -07:00 committed by TensorFlower Gardener
parent 96925bfe3d
commit e2c4396b59
1 changed files with 6 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
tensorflow/cc/saved_model/reader.cc
View File

@ -48,12 +48,12 @@ Status ReadSavedModel(const string& export_dir, SavedModel* saved_model_proto) {
export_dir); export_dir);
} }
Status FindMetaGraphDef(const SavedModel& saved_model_proto, Status FindMetaGraphDef(const std::unordered_set<string>& tags,
const std::unordered_set<string>& tags, SavedModel* saved_model_proto,
MetaGraphDef* meta_graph_def) { MetaGraphDef* meta_graph_def) {
LOG(INFO) << "Reading meta graph with tags { " << absl::StrJoin(tags, " ") LOG(INFO) << "Reading meta graph with tags { " << absl::StrJoin(tags, " ")
<< " }"; << " }";
for (const MetaGraphDef& graph_def : saved_model_proto.meta_graphs()) { for (MetaGraphDef& graph_def : *saved_model_proto->mutable_meta_graphs()) {
// Get tags from the graph_def. // Get tags from the graph_def.
std::unordered_set<string> graph_tags; std::unordered_set<string> graph_tags;
for (const string& tag : graph_def.meta_info_def().tags()) { for (const string& tag : graph_def.meta_info_def().tags()) {
@ -61,7 +61,7 @@ Status FindMetaGraphDef(const SavedModel& saved_model_proto,
} }
// Match with the set of tags provided. // Match with the set of tags provided.
if (graph_tags == tags) { if (graph_tags == tags) {
*meta_graph_def = graph_def; *meta_graph_def = std::move(graph_def);
return Status::OK(); return Status::OK();
} }
} }
@ -81,7 +81,8 @@ Status ReadMetaGraphDefFromSavedModel(const string& export_dir,
MetaGraphDef* const meta_graph_def) { MetaGraphDef* const meta_graph_def) {
SavedModel saved_model_proto; SavedModel saved_model_proto;
TF_RETURN_IF_ERROR(ReadSavedModel(export_dir, &saved_model_proto)); TF_RETURN_IF_ERROR(ReadSavedModel(export_dir, &saved_model_proto));
TF_RETURN_IF_ERROR(FindMetaGraphDef(saved_model_proto, tags, meta_graph_def)); TF_RETURN_IF_ERROR(
FindMetaGraphDef(tags, &saved_model_proto, meta_graph_def));
return Status::OK(); return Status::OK();
} }