Caffe2 - C++ API
A deep learning, cross platform ML framework
queue_ops.cc
1 
17 #include "queue_ops.h"
18 #include <memory>
19 #include "caffe2/utils/math.h"
20 
21 namespace caffe2 {
22 
23 CAFFE_KNOWN_TYPE(std::shared_ptr<BlobsQueue>);
24 
25 REGISTER_CPU_OPERATOR(CreateBlobsQueue, CreateBlobsQueueOp<CPUContext>);
26 REGISTER_CPU_OPERATOR(EnqueueBlobs, EnqueueBlobsOp<CPUContext>);
27 REGISTER_CPU_OPERATOR(DequeueBlobs, DequeueBlobsOp<CPUContext>);
28 REGISTER_CPU_OPERATOR(CloseBlobsQueue, CloseBlobsQueueOp<CPUContext>);
29 
30 REGISTER_CPU_OPERATOR(SafeEnqueueBlobs, SafeEnqueueBlobsOp<CPUContext>);
31 REGISTER_CPU_OPERATOR(SafeDequeueBlobs, SafeDequeueBlobsOp<CPUContext>);
32 REGISTER_CPU_OPERATOR(
33  WeightedSampleDequeueBlobs,
34  WeightedSampleDequeueBlobsOp<CPUContext>);
35 
36 OPERATOR_SCHEMA(CreateBlobsQueue).NumInputs(0).NumOutputs(1);
37 OPERATOR_SCHEMA(EnqueueBlobs)
38  .NumInputsOutputs([](int inputs, int outputs) {
39  return inputs >= 2 && outputs >= 1 && inputs == outputs + 1;
40  })
41  .EnforceInplace([](int input, int output) { return input == output + 1; });
42 OPERATOR_SCHEMA(DequeueBlobs)
43  .NumInputsOutputs([](int inputs, int outputs) {
44  return inputs == 1 && outputs >= 1;
45  })
46  .SetDoc(R"DOC(
47  Dequeue the blobs from queue.
48  )DOC")
49  .Arg("timeout_secs", "Timeout in secs, default: no timeout")
50  .Input(0, "queue", "The shared pointer for the BlobsQueue")
51  .Output(0, "blob", "The blob to store the dequeued data");
52 
53 OPERATOR_SCHEMA(CloseBlobsQueue).NumInputs(1).NumOutputs(0);
54 
55 OPERATOR_SCHEMA(SafeEnqueueBlobs)
56  .NumInputsOutputs([](int inputs, int outputs) {
57  return inputs >= 2 && outputs >= 2 && inputs == outputs;
58  })
59  .EnforceInplace([](int input, int output) { return input == output + 1; })
60  .SetDoc(R"DOC(
61 Enqueue the blobs into queue. When the queue is closed and full, the output
62 status will be set to true which can be used as exit criteria for execution
63 step.
64 The 1st input is the queue and the last output is the status. The rest are
65 data blobs.
66 )DOC")
67  .Input(0, "queue", "The shared pointer for the BlobsQueue");
68 
69 OPERATOR_SCHEMA(SafeDequeueBlobs)
70  .NumInputsOutputs([](int inputs, int outputs) {
71  return inputs == 1 && outputs >= 2;
72  })
73  .SetDoc(R"DOC(
74 Dequeue the blobs from queue. When the queue is closed and empty, the output
75 status will be set to true which can be used as exit criteria for execution
76 step.
77 The 1st input is the queue and the last output is the status. The rest are
78 data blobs.
79 )DOC")
80  .Arg(
81  "num_records",
82  "(default 1) If > 1, multiple records will be dequeued and tensors "
83  "for each column will be concatenated. This requires all tensors in "
84  "the records to be at least 1D, and to have the same inner dimensions.")
85  .Input(0, "queue", "The shared pointer for the BlobsQueue")
86  .Output(0, "blob", "The blob to store the dequeued data")
87  .Output(1, "status", "Is set to 0/1 depending on the success of dequeue");
88 
89 OPERATOR_SCHEMA(WeightedSampleDequeueBlobs)
90  .NumInputs(1, INT_MAX)
91  .NumOutputs(2, INT_MAX)
92  .SetDoc(R"DOC(
93 Dequeue the blobs from multiple queues. When one of queues is closed and empty,
94 the output status will be set to true which can be used as exit criteria for
95 execution step.
96 The 1st input is the queue and the last output is the status. The rest are
97 data blobs.
98 )DOC")
99  .Arg("weights", "Weights for sampling from multiple queues")
100  .Arg(
101  "table_idx_blob",
102  "The index of the blob (among the output blob list) "
103  "that will be used to store the index of the table chosen to read the "
104  "current batch.");
105 
106 NO_GRADIENT(CreateBlobsQueue);
107 NO_GRADIENT(EnqueueBlobs);
108 NO_GRADIENT(DequeueBlobs);
109 NO_GRADIENT(CloseBlobsQueue);
110 
111 NO_GRADIENT(SafeEnqueueBlobs);
112 NO_GRADIENT(SafeDequeueBlobs);
113 NO_GRADIENT(WeightedSampleDequeueBlobs);
114 
115 }
Copyright (c) 2016-present, Facebook, Inc.