1 /* Copyright 2016 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 16 /// SavedModel loading functions and SavedModelBundle struct. 17 18 #ifndef TENSORFLOW_CC_SAVED_MODEL_LOADER_H_ 19 #define TENSORFLOW_CC_SAVED_MODEL_LOADER_H_ 20 21 #include <string> 22 #include <unordered_set> 23 24 #include "tensorflow/core/lib/core/status.h" 25 #include "tensorflow/core/protobuf/graph_debug_info.pb.h" 26 #include "tensorflow/core/protobuf/meta_graph.pb.h" 27 #include "tensorflow/core/public/session.h" 28 29 namespace tensorflow { 30 31 /// Represents a SavedModel that is loaded from storage. 32 class SavedModelBundleInterface { 33 public: 34 virtual ~SavedModelBundleInterface(); 35 36 /// Returns the TensorFlow Session that can be used to interact with the 37 /// SavedModel. 38 virtual Session* GetSession() const = 0; 39 40 /// Returns a map from signature name to SignatureDef for all signatures in 41 /// in the SavedModel. 42 virtual const protobuf::Map<string, SignatureDef>& GetSignatures() const = 0; 43 }; 44 45 /// SavedModel representation once the SavedModel is loaded from storage. 46 /// 47 /// NOTE: Prefer to use SavedModelBundleLite in new code, as it consumes less 48 /// RAM. 49 struct SavedModelBundle : public SavedModelBundleInterface { 50 /// A TensorFlow Session does not Close itself on destruction. To avoid 51 /// resource leaks, we explicitly call Close on Sessions that we create. ~SavedModelBundleSavedModelBundle52 ~SavedModelBundle() override { 53 if (session) { 54 session->Close().IgnoreError(); 55 } 56 } 57 58 SavedModelBundle() = default; 59 GetSessionSavedModelBundle60 Session* GetSession() const override { return session.get(); } GetSignaturesSavedModelBundle61 const protobuf::Map<string, SignatureDef>& GetSignatures() const override { 62 return meta_graph_def.signature_def(); 63 } 64 65 std::unique_ptr<Session> session; 66 MetaGraphDef meta_graph_def; 67 std::unique_ptr<GraphDebugInfo> debug_info; 68 }; 69 70 // A version of SavedModelBundle that avoids storing a potentially large 71 // MetaGraphDef. Prefer to use SavedModelBundleLite in new code. 72 class SavedModelBundleLite : public SavedModelBundleInterface { 73 public: 74 SavedModelBundleLite() = default; 75 SavedModelBundleLite& operator=(SavedModelBundleLite&& other) = default; 76 SavedModelBundleLite(std::unique_ptr<Session> session,protobuf::Map<string,SignatureDef> signatures)77 SavedModelBundleLite(std::unique_ptr<Session> session, 78 protobuf::Map<string, SignatureDef> signatures) 79 : session_(std::move(session)), signatures_(std::move(signatures)) {} 80 81 /// A TensorFlow Session does not Close itself on destruction. To avoid 82 /// resource leaks, we explicitly call Close on Sessions that we create. ~SavedModelBundleLite()83 ~SavedModelBundleLite() override { 84 if (session_) { 85 session_->Close().IgnoreError(); 86 } 87 } 88 GetSession()89 Session* GetSession() const override { return session_.get(); } GetSignatures()90 const protobuf::Map<string, SignatureDef>& GetSignatures() const override { 91 return signatures_; 92 } 93 94 private: 95 std::unique_ptr<Session> session_; 96 protobuf::Map<string, SignatureDef> signatures_; 97 }; 98 99 // Restore variable and resources in the SavedModel export dir for the 100 // indicated metagraph. 101 // The recommended way to load a saved model is to call LoadSavedModel, 102 // which provides an already initialized Metagraph, Session, and DebugInfo. 103 Status RestoreSession(const RunOptions& run_options, 104 const MetaGraphDef& meta_graph, const string& export_dir, 105 std::unique_ptr<Session>* session); 106 107 // Initialize a session which wraps this metagraph. 108 // The recommended way to load a saved model is to call LoadSavedModel, 109 // which provides an already initialized Metagraph, Session, and DebugInfo. 110 Status LoadMetagraphIntoSession(const SessionOptions& session_options, 111 const MetaGraphDef& meta_graph, 112 std::unique_ptr<Session>* session); 113 114 /// Loads a SavedModel from the specified export directory. The MetaGraphDef 115 /// to be loaded is identified by the supplied tags, corresponding exactly to 116 /// the set of tags used at SavedModel build time. Stores a SavedModel bundle in 117 /// *bundle with a session and the requested MetaGraphDef, if found. 118 /// 119 /// NOTE: Prefer the overload that takes a SavedModelBundleLite* in new code. 120 Status LoadSavedModel(const SessionOptions& session_options, 121 const RunOptions& run_options, const string& export_dir, 122 const std::unordered_set<string>& tags, 123 SavedModelBundle* const bundle); 124 125 /// Loads a SavedModel from the specified export directory. The MetaGraphDef 126 /// to be loaded is identified by the supplied tags, corresponding exactly to 127 /// the set of tags used at SavedModel build time. Stores a SavedModel bundle 128 /// in *bundle with a session created from the requested MetaGraphDef if found. 129 /// 130 /// This overload creates a SavedModelBundleLite, which consumes less RAM than 131 /// an equivalent SavedModelBundle. 132 Status LoadSavedModel(const SessionOptions& session_options, 133 const RunOptions& run_options, const string& export_dir, 134 const std::unordered_set<string>& tags, 135 SavedModelBundleLite* const bundle); 136 137 /// Checks whether the provided directory could contain a SavedModel. Note that 138 /// the method does not load any data by itself. If the method returns `false`, 139 /// the export directory definitely does not contain a SavedModel. If the method 140 /// returns `true`, the export directory may contain a SavedModel but provides 141 /// no guarantee that it can be loaded. 142 bool MaybeSavedModelDirectory(const std::string& export_dir); 143 144 } // namespace tensorflow 145 146 #endif // TENSORFLOW_CC_SAVED_MODEL_LOADER_H_ 147