Caffe2 - C++ API
A deep learning, cross platform ML framework
load_save_op.cc
1 #include "caffe2/operators/load_save_op.h"
2 
3 namespace caffe2 {
4 
5 template <>
6 void LoadOp<CPUContext>::SetCurrentDevice(BlobProto* proto) {
7  if (proto->has_tensor()) {
8  proto->mutable_tensor()->clear_device_detail();
9  proto->mutable_tensor()->mutable_device_detail()->set_device_type(
10  PROTO_CPU);
11  }
12 }
13 
14 REGISTER_CPU_OPERATOR(DBExists, DBExistsOp<CPUContext>);
15 REGISTER_CPU_OPERATOR(Load, LoadOp<CPUContext>);
16 REGISTER_CPU_OPERATOR(Save, SaveOp<CPUContext>);
17 REGISTER_CPU_OPERATOR(Checkpoint, CheckpointOp<CPUContext>);
18 // CPU Operator old name: do NOT use, we may deprecate this later.
19 REGISTER_CPU_OPERATOR(Snapshot, CheckpointOp<CPUContext>);
20 
21 OPERATOR_SCHEMA(DBExists)
22  .NumInputs(0)
23  .NumOutputs(1)
24  .SetDoc(R"DOC(
25 Checks if the db described by the arguments exists.
26 
27 Github Links:
28 
29 - https://github.com/pytorch/pytorch/blob/master/caffe2/operators/load_save_op.cc
30 
31 <details>
32 
33 <summary> <b>Example</b> </summary>
34 
35 **Code**
36 
37 ```
38 workspace.ResetWorkspace()
39 
40 op = core.CreateOperator(
41  "DBExists",
42  [],
43  ["exists"],
44  db_name="test_db",
45  db_type="leveldb",
46 )
47 
48 workspace.RunOperatorOnce(op)
49 print("exists:", workspace.FetchBlob("exists"))
50 
51 ```
52 
53 </details>
54 
55 )DOC")
56  .Output(0, "exists", "*(type: Tensor`<bool>`)* Scalar boolean output "
57  "tensor. True if the db exists, else false.")
58  .Arg(
59  "absolute_path",
60  "*(type: int; default: 0)* If set to non-zero, save the db directly to "
61  "the path specified by the `db` arg. If not set (default), prepend the "
62  "path of the current root folder of the workspace to the path specified "
63  "by the `db` arg.")
64  .Arg("db_name", "*(type: string)* Path to the db in question; see the "
65  "`absolute_path` arg details for options regarding the current root folder "
66  "of the workspace.")
67  .Arg("db_type", "*(type: string)* Type of db to save (options: \"lmdb\", "
68  "\"leveldb\", \"minidb\").");
69 
70 OPERATOR_SCHEMA(Load)
71  .NumInputs(0, INT_MAX)
72  .NumOutputs(0, INT_MAX)
73  .SetDoc(R"DOC(
74 The Load operator loads a set of serialized blobs from a db or multiple dbs. It
75 takes $[0, \infty)$ number of inputs and $[0, \infty)$ number of outputs, using
76 the db keys to match the db entries with the outputs.
77 
78 If at least one input is passed, then it is assumed that that input blobs are a
79 set of DBReaders to load from. Otherwise the `db` or `dbs` argument is used to load
80 blobs from one single db or multiple dbs respectively. `db_type` argument is used
81 to specify the type of the input db/dbs.
82 
83 Github Links:
84 
85 - https://github.com/pytorch/pytorch/blob/master/caffe2/operators/load_save_op.cc
86 
87 <details>
88 
89 <summary> <b>Example</b> </summary>
90 
91 **Code**
92 
93 ```
94 workspace.ResetWorkspace()
95 
96 op = core.CreateOperator(
97  "Load",
98  [],
99  ["X", "Y"],
100  db="test_db",
101  db_type="lmdb"
102 )
103 
104 workspace.RunOperatorOnce(op)
105 print("X:", workspace.FetchBlob("X"))
106 print("Y:", workspace.FetchBlob("Y"))
107 
108 ```
109 
110 </details>
111 
112 )DOC")
113  .Input(
114  0,
115  "X, Y, ...",
116  "*(type: List(DBReader))* [OPTIONAL] List of DBReaders to load from. Can "
117  "use this instead of the `db`/`dbs` args.")
118  .Arg(
119  "absolute_path",
120  "*(type: int; default: 0)* If set to non-zero, save the db directly to "
121  "the path specified by the `db` arg. If not set (default), prepend the "
122  "path of the current root folder of the workspace to the path specified "
123  "by the `db` arg.")
124  .Arg(
125  "add_prefix",
126  "*(type: string, default: \"\")* Blobs will be prefixed with this when "
127  "loading. Useful for avoiding collisions with blobs existing in the "
128  "workspace. The output blob names specified to this op should include "
129  "this prefix.")
130  .Arg(
131  "strip_prefix",
132  "*(type: string, default: \"\")* Characters in the provided blob names "
133  "that match `strip_prefix` will be removed prior to saving. Also, "
134  "characters that precede `strip_prefix` will be removed. Useful for "
135  "removing device scope from blob names.")
136  .Arg("db", "*(type: string)* The output path of the db. See the "
137  "`absolute_path` arg details for options regarding the current root folder "
138  "of the workspace.")
139  .Arg(
140  "dbs",
141  "*(type: List(string))* List of paths to dbs to load blobs from. See "
142  "the `absolute_path` arg details for options regarding the current "
143  "root folder of the workspace.")
144  .Arg("db_type", "(type: string)* Type of db to save (options: \"lmdb\", "
145  "\"leveldb\", \"minidb\").")
146  .Arg(
147  "keep_device",
148  "*(type: int; default: 0)* If nonzero, the blobs are loaded into the "
149  "device that is specified in the serialized `BlobProto`. Otherwise, "
150  "the device will be set as the one that the `Load` operator is being "
151  "run under.")
152  .Arg(
153  "load_all",
154  "*(type: int; default: 0)* If nonzero, will load all blobs pointed to "
155  "by the db to the workspace overwriting/creating blobs as needed.")
156  .Arg(
157  "allow_incomplete",
158  "*(type: bool; default: False)* If True, will allow not loading all "
159  "the output blobs specified in the outputs.")
160  .Arg(
161  "source_blob_names",
162  "*(type: List(string))* If set, used instead of output blob names to "
163  "specify which blobs in the db shall be loaded. Must be the same "
164  "length as number of output blobs.");
165 
166 OPERATOR_SCHEMA(Save)
167  .NumInputs(1, INT_MAX)
168  .NumOutputs(0)
169  .SetDoc(R"DOC(
170 Saves a set of blobs to a db. It takes $[1, \infty)$ number of inputs and has
171 no output. The contents of the inputs are written into the db using the
172 settings specified by the arguments.
173 
174 Github Links:
175 
176 - https://github.com/pytorch/pytorch/blob/master/caffe2/operators/load_save_op.cc
177 
178 <details>
179 
180 <summary> <b>Example</b> </summary>
181 
182 **Code**
183 
184 ```
185 workspace.ResetWorkspace()
186 
187 op = core.CreateOperator(
188  "Save",
189  ["X", "Y", "Z"],
190  [],
191  db="test_db2",
192  db_type="leveldb",
193  blob_name_overrides=["x_scores", "y_scores", "z_scores"]
194 )
195 
196 workspace.FeedBlob("X", np.random.randint(20, size=(5,5)))
197 workspace.FeedBlob("Y", np.random.randint(20, size=(5,5)))
198 workspace.FeedBlob("Z", np.random.randint(20, size=(5,5)))
199 workspace.RunOperatorOnce(op)
200 
201 ```
202 
203 </details>
204 
205 )DOC")
206  .Arg(
207  "absolute_path",
208  "*(type: int; default: 0)* If set to non-zero, save the db directly to "
209  "the path specified by the `db` arg. If not set (default), prepend the "
210  "path of the current root folder of the workspace to the path specified "
211  "by the `db` arg.")
212  .Arg(
213  "strip_prefix",
214  "*(type: string, default: \"\")* Characters in the provided blob names "
215  "that match `strip_prefix` will be removed prior to saving. Also, "
216  "characters that precede `strip_prefix` will be removed. Useful for "
217  "removing device scope from blob names.")
218  .Arg(
219  "blob_name_overrides",
220  "*(List(string))* If set, used as blob names instead of original blob "
221  "names. Must be same length as number of blobs.")
222  .Arg("db", "*(type: string)* The output path of the db. See the "
223  "`absolute_path` arg details for options regarding the current root folder "
224  "of the workspace.")
225  .Arg("db_type", "*(type: string)* Type of db to save (options: \"lmdb\", "
226  "\"leveldb\", \"minidb\").")
227  .Arg("chunk_size", "*(type: string; default: kDefaultChunkSize)* The chunk "
228  "size to split tensor data into. If not set, caffe2_tensor_chunk_size will "
229  "be used")
230  .Input(0, "X", "*(type: Tensor)* Input tensor(s).");
231 
232 OPERATOR_SCHEMA(Checkpoint)
233  .NumInputs(1, INT_MAX)
234  .NumOutputs(0)
235  .SetDoc(R"DOC(
236 The Checkpoint operator is similar to the Save operator, but allows one to save
237 to db every few iterations, with a db name that is appended with the iteration
238 count. It takes [1, infinity) number of inputs and has no output. The first
239 input has to be a TensorCPU of type int and has size 1 (i.e. the iteration
240 counter). This is determined whether we need to do checkpointing.
241 )DOC")
242  .Arg(
243  "absolute_path",
244  "(int, default 0) if set, use the db path directly and do not prepend "
245  "the current root folder of the workspace.")
246  .Arg(
247  "db",
248  "(string) a template string that one can combine with the "
249  "iteration to create the final db name. For example, "
250  "\"/home/lonestarr/checkpoint_%08d.db\"")
251  .Arg("db_type", "(string) the type of the db.")
252  .Arg(
253  "every",
254  "(int, default 1) the checkpointing is carried out when "
255  "(iter mod every) is zero.");
256 
257 OPERATOR_SCHEMA(Snapshot);
258 
259 NO_GRADIENT(Load);
260 SHOULD_NOT_DO_GRADIENT(DBExists);
261 SHOULD_NOT_DO_GRADIENT(Save);
262 SHOULD_NOT_DO_GRADIENT(Checkpoint);
263 SHOULD_NOT_DO_GRADIENT(Snapshot);
264 } // namespace caffe2
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...
Definition: blob.h:13