1 /* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15 #include "tensorflow/cc/experimental/libexport/load.h"
16
17 #include "tensorflow/cc/saved_model/constants.h"
18 #include "tensorflow/core/framework/graph.pb.h"
19 #include "tensorflow/core/framework/node_def.pb.h"
20 #include "tensorflow/core/platform/env.h"
21 #include "tensorflow/core/platform/errors.h"
22 #include "tensorflow/core/platform/path.h"
23 #include "tensorflow/core/protobuf/meta_graph.pb.h"
24 #include "tensorflow/core/util/tensor_bundle/tensor_bundle.h"
25
26 #define RETURN_IF_ERROR(s) \
27 { \
28 auto c = (s); \
29 if (!c.ok()) return c; \
30 }
31
32 namespace tensorflow {
33 namespace libexport {
34
35 using protobuf::RepeatedPtrField;
36
Load(const std::string & path)37 tensorflow::StatusOr<TFPackage> TFPackage::Load(const std::string& path) {
38 // Load the proto
39 TFPackage tf_package;
40 const string saved_model_pb_path = io::JoinPath(path, kSavedModelFilenamePb);
41 const string saved_model_pbtxt_path =
42 io::JoinPath(path, kSavedModelFilenamePbTxt);
43 if (Env::Default()->FileExists(saved_model_pb_path).ok()) {
44 RETURN_IF_ERROR(ReadBinaryProto(Env::Default(), saved_model_pb_path,
45 &tf_package.saved_model_proto_));
46 } else if (Env::Default()->FileExists(saved_model_pbtxt_path).ok()) {
47 RETURN_IF_ERROR(ReadTextProto(Env::Default(), saved_model_pbtxt_path,
48 &tf_package.saved_model_proto_));
49 } else {
50 return Status(error::Code::NOT_FOUND,
51 "Could not find SavedModel .pb or .pbtxt at supplied export "
52 "directory path: " +
53 path);
54 }
55
56 // Load the trackable object graph for restoring checkpoint values
57 const std::string variables_dir =
58 tensorflow::io::JoinPath(path, tensorflow::kSavedModelVariablesDirectory);
59 // TODO(b/228181641): revisit non-explicit-checkpoint-loading behavior when
60 // MLAs come along
61 if (Env::Default()->FileExists(variables_dir).ok()) {
62 tf_package.has_checkpoint_ = true;
63 tf_package.variables_filepath_ = tensorflow::io::JoinPath(
64 variables_dir, tensorflow::kSavedModelVariablesFilename);
65 tf_package.variable_reader_ = std::make_unique<tensorflow::BundleReader>(
66 tensorflow::Env::Default(), tf_package.variables_filepath_);
67 tensorflow::Tensor object_graph_tensor;
68 RETURN_IF_ERROR(tf_package.variable_reader_->Lookup(
69 tensorflow::kObjectGraphProtoKey, &object_graph_tensor));
70 const auto* object_graph_string =
71 reinterpret_cast<const tensorflow::tstring*>(
72 object_graph_tensor.tensor_data().data());
73 // TODO(danielellis): make sure parse was successful
74 tf_package.trackable_object_graph_.ParseFromString(*object_graph_string);
75 } else {
76 tf_package.has_checkpoint_ = false;
77 LOG(INFO)
78 << "No checkpoint found, assuming this is a program-only SavedModel";
79 }
80
81 // Build a map of node names to their corresponding nodes.
82 //
83 // See `GetGraphDefNode` for more details.
84 const auto& nodes =
85 tf_package.saved_model_proto_.meta_graphs(0).graph_def().node();
86 for (const auto& node : nodes) {
87 tf_package.graph_def_nodes_by_name_[node.name()] = &node;
88 }
89 return tf_package;
90 }
91
GetVariableCheckpointKey(int index)92 tensorflow::StatusOr<std::string> TFPackage::GetVariableCheckpointKey(
93 int index) {
94 // TODO(danielellis): make sure valid index
95 const auto& trackable_object = trackable_object_graph_.nodes(index);
96 const TrackableObjectGraph::TrackableObject::SerializedTensor*
97 serialized_tensor = nullptr;
98 for (auto& maybe_serialized_tensor : trackable_object.attributes()) {
99 if (maybe_serialized_tensor.name() == "VARIABLE_VALUE") {
100 serialized_tensor = &maybe_serialized_tensor;
101 }
102 }
103 if (serialized_tensor == nullptr) {
104 return tensorflow::Status(error::INTERNAL,
105 "Failed to find variable value field.");
106 }
107 return serialized_tensor->checkpoint_key();
108 }
109
GetObjectGraph()110 const SavedObjectGraph& TFPackage::GetObjectGraph() {
111 return saved_model_proto_.mutable_meta_graphs(0)->object_graph_def();
112 }
113
GetGraphDefNode(std::string name)114 tensorflow::StatusOr<const tensorflow::NodeDef*> TFPackage::GetGraphDefNode(
115 std::string name) {
116 const auto& iter = graph_def_nodes_by_name_.find(name);
117 if (iter == graph_def_nodes_by_name_.end()) {
118 return tensorflow::Status(error::INTERNAL,
119 absl::StrCat("Failed to find node named ", name));
120 }
121 return iter->second;
122 }
123
GetFunctionDefs()124 const RepeatedPtrField<FunctionDef>& TFPackage::GetFunctionDefs() {
125 auto& function_library =
126 saved_model_proto_.mutable_meta_graphs(0)->graph_def().library();
127 return function_library.function();
128 }
129
130 } // namespace libexport
131 } // namespace tensorflow
132