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