5 constexpr
auto kBlobName =
"blob_name";
6 constexpr
auto kAddValue =
"add_value";
8 StoreSetOp::StoreSetOp(
const OperatorDef& operator_def, Workspace* ws)
9 : Operator<CPUContext>(operator_def, ws),
11 GetSingleArgument<
std::string>(kBlobName, operator_def.input(DATA))) {
14 bool StoreSetOp::RunOnDevice() {
17 OperatorBase::Input<std::unique_ptr<StoreHandler>>(HANDLER).
get();
18 handler->set(blobName_, InputBlob(DATA).Serialize(blobName_));
22 REGISTER_CPU_OPERATOR(StoreSet, StoreSetOp);
23 OPERATOR_SCHEMA(StoreSet)
27 Set a blob in a store. The key is the input blob's name and the value 28 is the data in that blob. The key can be overridden by specifying the 31 .Arg("blob_name",
"alternative key for the blob (optional)")
32 .Input(0,
"handler",
"unique_ptr<StoreHandler>")
33 .Input(1,
"data",
"data blob");
35 StoreGetOp::StoreGetOp(
const OperatorDef& operator_def, Workspace* ws)
36 : Operator<CPUContext>(operator_def, ws),
37 blobName_(GetSingleArgument<
std::string>(
39 operator_def.output(DATA))) {}
41 bool StoreGetOp::RunOnDevice() {
44 OperatorBase::Input<std::unique_ptr<StoreHandler>>(HANDLER).
get();
45 OperatorBase::Outputs()[DATA]->Deserialize(handler->get(blobName_));
49 REGISTER_CPU_OPERATOR(StoreGet, StoreGetOp);
50 OPERATOR_SCHEMA(StoreGet)
54 Get a blob from a store. The key is the output blob's name. The key 55 can be overridden by specifying the 'blob_name' argument. 57 .Arg("blob_name",
"alternative key for the blob (optional)")
58 .Input(0,
"handler",
"unique_ptr<StoreHandler>")
59 .Output(0,
"data",
"data blob");
61 StoreAddOp::StoreAddOp(
const OperatorDef& operator_def, Workspace* ws)
62 : Operator<CPUContext>(operator_def, ws),
63 blobName_(GetSingleArgument<
std::string>(kBlobName,
"")),
64 addValue_(GetSingleArgument<int64_t>(kAddValue, 1)) {
65 CAFFE_ENFORCE(HasArgument(kBlobName));
68 bool StoreAddOp::RunOnDevice() {
70 OperatorBase::Input<std::unique_ptr<StoreHandler>>(HANDLER).
get();
71 Output(VALUE)->Resize(1);
72 Output(VALUE)->mutable_data<int64_t>()[0] =
73 handler->add(blobName_, addValue_);
77 REGISTER_CPU_OPERATOR(StoreAdd, StoreAddOp);
78 OPERATOR_SCHEMA(StoreAdd)
82 Add a value to a remote counter. If the key is not set, the store 83 initializes it to 0 and then performs the add operation. The operation 84 returns the resulting counter value. 86 .Arg("blob_name",
"key of the counter (required)")
87 .Arg(
"add_value",
"value that is added (optional, default: 1)")
88 .Input(0,
"handler",
"unique_ptr<StoreHandler>")
89 .Output(0,
"value",
"the current value of the counter");
91 StoreWaitOp::StoreWaitOp(
const OperatorDef& operator_def, Workspace* ws)
92 : Operator<CPUContext>(operator_def, ws),
93 blobNames_(GetRepeatedArgument<
std::string>(kBlobName)) {}
95 bool StoreWaitOp::RunOnDevice() {
97 OperatorBase::Input<std::unique_ptr<StoreHandler>>(HANDLER).
get();
98 if (InputSize() == 2 && Input(1).IsType<std::string>()) {
100 blobNames_.empty(),
"cannot specify both argument and input blob");
101 std::vector<std::string> blobNames;
102 auto* namesPtr = Input(1).data<std::string>();
103 for (
int i = 0; i < Input(1).size(); ++i) {
104 blobNames.push_back(namesPtr[i]);
106 handler->wait(blobNames);
108 handler->wait(blobNames_);
113 REGISTER_CPU_OPERATOR(StoreWait, StoreWaitOp);
114 OPERATOR_SCHEMA(StoreWait)
118 Wait for the specified blob names to be set. The blob names can be passed 119 either as an input blob with blob names or as an argument. 121 .Arg("blob_names",
"names of the blobs to wait for (optional)")
122 .Input(0,
"handler",
"unique_ptr<StoreHandler>")
123 .Input(1,
"names",
"names of the blobs to wait for (optional)");
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...