Caffe2 - C++ API
A deep learning, cross platform ML framework
load_save_op.cc
1 
17 #include "caffe2/operators/load_save_op.h"
18 
19 namespace caffe2 {
20 
21 template <>
22 void LoadOp<CPUContext>::SetCurrentDevice(BlobProto* proto) {
23  if (proto->has_tensor()) {
24  proto->mutable_tensor()->mutable_device_detail()->set_device_type(CPU);
25  }
26 }
27 
28 REGISTER_CPU_OPERATOR(DBExists, DBExistsOp<CPUContext>);
29 REGISTER_CPU_OPERATOR(Load, LoadOp<CPUContext>);
30 REGISTER_CPU_OPERATOR(Save, SaveOp<CPUContext>);
31 REGISTER_CPU_OPERATOR(Checkpoint, CheckpointOp<CPUContext>);
32 // CPU Operator old name: do NOT use, we may deprecate this later.
33 REGISTER_CPU_OPERATOR(Snapshot, CheckpointOp<CPUContext>);
34 
35 OPERATOR_SCHEMA(DBExists)
36  .NumInputs(0)
37  .NumOutputs(1)
38  .SetDoc(R"DOC(
39 Checks if the DB exists.
40 )DOC")
41  .Output(0, "exists", "A scalar bool Tensor.")
42  .Arg(
43  "absolute_path",
44  "(int, default 0) if set, use the db path directly and do not prepend "
45  "the current root folder of the workspace.")
46  .Arg("db_name", "(string) the path to the db to load.")
47  .Arg("db_type", "(string) the type of the db.");
48 
49 OPERATOR_SCHEMA(Load)
50  .NumInputs(0, INT_MAX)
51  .NumOutputs(0, INT_MAX)
52  .SetDoc(R"DOC(
53 The Load operator loads a set of serialized blobs from a db or multiple dbs. It
54 takes [0, infinity) number of inputs and [0, infinity) number of outputs, using
55 the db keys to match the db entries with the outputs.
56 
57 If at least one input is passed, then it is assumed that that input blobs are a
58 set of DBReaders to load from. Otherwise the db or dbs argument is used to load
59 blobs from one single db or multiple dbs respectively. db_type argument is used
60 to specify the type of the input db/dbs.
61 )DOC")
62  .Arg(
63  "absolute_path",
64  "(int, default 0) if set, use the db path directly and do not prepend "
65  "the current root folder of the workspace.")
66  .Arg(
67  "add_prefix",
68  "(string, default=\"\") blobs will be prefixed with this when loading."
69  "Useful for avoiding collisions with blobs existing in the workspace."
70  "The output blob names specified to this op should include this prefix.")
71  .Arg(
72  "strip_prefix",
73  "(string, default=\"\") characters in the provided blob "
74  " names that match strip_prefix will be removed prior to loading."
75  " Also, characters that precede strip_prefix will be removed. Useful "
76  " for removing device scope from blob names.")
77  .Arg("db", "(string) the path to the db to load.")
78  .Arg(
79  "dbs",
80  "(list of strings) the paths to the dbs to load. This is used for loading"
81  " blobs from multiple databases. If it is set, argument in \"db\" will be"
82  " ignored.")
83  .Arg("db_type", "(string) the type of the db.")
84  .Arg(
85  "keep_device",
86  "(int, default 0) if nonzero, the blobs are loaded into the device that "
87  "is specified in the serialized BlobProto. Otherwise, the device will be "
88  "set as the one that the Load operator is being run under.")
89  .Arg(
90  "load_all",
91  "(int, default 0) if nonzero, will load all blobs pointed to by the db "
92  "to the workspace overwriting/creating blobs as needed.")
93  .Arg(
94  "allow_incomplete",
95  "(bool, default false) if true, will allow not loading all the output "
96  "blobs specified in the outputs")
97  .Arg(
98  "source_blob_names",
99  "(list of strings) if set, used instead of output "
100  "blob names, to specify which blobs in the db shall be loaded. Must be "
101  "the same length as number of output blobs.");
102 
103 OPERATOR_SCHEMA(Save)
104  .NumInputs(1, INT_MAX)
105  .NumOutputs(0)
106  .SetDoc(R"DOC(
107 The Save operator saves a set of blobs to a db. It takes [1, infinity) number
108 of inputs and has no output. The contents of the inputs are written into the
109 db specified by the arguments.
110 )DOC")
111  .Arg(
112  "absolute_path",
113  "(int, default 0) if set, use the db path directly and do not prepend "
114  "the current root folder of the workspace.")
115  .Arg(
116  "strip_prefix",
117  "(string, default=\"\") characters in the provided blob "
118  " names that match strip_prefix will be removed prior to saving."
119  " Also, characters that precede strip_prefix will be removed. Useful "
120  " for removing device scope from blob names.")
121  .Arg(
122  "blob_name_overrides",
123  "(list of strings) if set, used instead of original "
124  "blob names. Must be the same length as number of blobs.")
125  .Arg("db", "(string) the path to the db to load.")
126  .Arg("db_type", "(string) the type of the db.");
127 
128 OPERATOR_SCHEMA(Checkpoint)
129  .NumInputs(1, INT_MAX)
130  .NumOutputs(0)
131  .SetDoc(R"DOC(
132 The Checkpoint operator is similar to the Save operator, but allows one to save
133 to db every few iterations, with a db name that is appended with the iteration
134 count. It takes [1, infinity) number of inputs and has no output. The first
135 input has to be a TensorCPU of type int and has size 1 (i.e. the iteration
136 counter). This is determined whether we need to do checkpointing.
137 )DOC")
138  .Arg(
139  "absolute_path",
140  "(int, default 0) if set, use the db path directly and do not prepend "
141  "the current root folder of the workspace.")
142  .Arg(
143  "db",
144  "(string) a template string that one can combine with the "
145  "iteration to create the final db name. For example, "
146  "\"/home/lonestarr/checkpoint_%08d.db\"")
147  .Arg("db_type", "(string) the type of the db.")
148  .Arg(
149  "every",
150  "(int, default 1) the checkpointing is carried out when "
151  "(iter mod every) is zero.");
152 
153 OPERATOR_SCHEMA(Snapshot);
154 
155 NO_GRADIENT(Load);
156 SHOULD_NOT_DO_GRADIENT(DBExists);
157 SHOULD_NOT_DO_GRADIENT(Save);
158 SHOULD_NOT_DO_GRADIENT(Checkpoint);
159 SHOULD_NOT_DO_GRADIENT(Snapshot);
160 } // namespace caffe2
Copyright (c) 2016-present, Facebook, Inc.