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 /// Loads a SavedModel from the specified export directory. The MetaGraphDef 100 /// to be loaded is identified by the supplied tags, corresponding exactly to 101 /// the set of tags used at SavedModel build time. Stores a SavedModel bundle in 102 /// *bundle with a session and the requested MetaGraphDef, if found. 103 /// 104 /// NOTE: Prefer the overload that takes a SavedModelBundleLite* in new code. 105 Status LoadSavedModel(const SessionOptions& session_options, 106 const RunOptions& run_options, const string& export_dir, 107 const std::unordered_set<string>& tags, 108 SavedModelBundle* const bundle); 109 110 /// Loads a SavedModel from the specified export directory. The MetaGraphDef 111 /// to be loaded is identified by the supplied tags, corresponding exactly to 112 /// the set of tags used at SavedModel build time. Stores a SavedModel bundle 113 /// in *bundle with a session created from the requested MetaGraphDef if found. 114 /// 115 /// This overload creates a SavedModelBundleLite, which consumes less RAM than 116 /// an equivalent SavedModelBundle. 117 Status LoadSavedModel(const SessionOptions& session_options, 118 const RunOptions& run_options, const string& export_dir, 119 const std::unordered_set<string>& tags, 120 SavedModelBundleLite* const bundle); 121 122 /// Checks whether the provided directory could contain a SavedModel. Note that 123 /// the method does not load any data by itself. If the method returns `false`, 124 /// the export directory definitely does not contain a SavedModel. If the method 125 /// returns `true`, the export directory may contain a SavedModel but provides 126 /// no guarantee that it can be loaded. 127 bool MaybeSavedModelDirectory(const string& export_dir); 128 129 } // namespace tensorflow 130 131 #endif // TENSORFLOW_CC_SAVED_MODEL_LOADER_H_ 132