1 #ifndef CAFFE2_OPERATORS_BATCH_GATHER_OPS_H_ 2 #define CAFFE2_OPERATORS_BATCH_GATHER_OPS_H_ 4 #include "caffe2/core/context.h" 5 #include "caffe2/core/operator.h" 6 #include "caffe2/utils/math.h" 10 template <
class Context>
13 USE_OPERATOR_CONTEXT_FUNCTIONS;
16 bool RunOnDevice()
override {
18 this, OperatorBase::Input<TensorCPU>(INDICES));
21 template <
typename TInd>
22 bool DoRunWithType() {
23 auto& data = Input(DATA);
24 auto& indices = Input(INDICES);
25 auto* output = Output(0);
27 CAFFE_ENFORCE_GE(data.ndim(), 2,
"DATA should be at least 2-D");
30 shape.push_back(data.dim(0));
31 shape.insert(shape.end(), indices.dims().begin(), indices.dims().end());
32 shape.insert(shape.end(), data.dims().begin() + 2, data.dims().end());
33 output->Resize(shape);
35 auto block_size = data.size_from_dim(2);
36 auto block_bytesize = block_size * data.meta().itemsize();
37 auto N = indices.size();
38 auto data_batch_bytesize = data.size_from_dim(1) * data.meta().itemsize();
39 auto gathered_batch_bytesize =
40 N * data.size_from_dim(2) * data.meta().itemsize();
41 const TInd* idxs = indices.template data<TInd>();
42 auto src_base =
static_cast<const char*
>(data.raw_data());
43 auto out =
static_cast<char*
>(output->raw_mutable_data(data.meta()));
45 for (
auto batch = 0; batch < data.dim(0); ++batch) {
46 for (
auto i = 0; i < N; ++i) {
49 0 <= idx && idx < data.dim(1),
50 "INDICES element is out of DATA bounds, id=",
55 src_base + idx * block_bytesize + batch * data_batch_bytesize;
56 auto dst = out + i * block_bytesize + batch * gathered_batch_bytesize;
57 context_.template CopyItems<Context, Context>(
58 data.meta(), block_size, src, dst);
64 INPUT_TAGS(DATA, INDICES);
67 template <
class Context>
70 USE_OPERATOR_CONTEXT_FUNCTIONS;
73 bool RunOnDevice()
override {
75 this, OperatorBase::Input<TensorCPU>(INDICES));
78 template <
typename TInd>
79 bool DoRunWithType() {
82 TInd>::call(
this, Input(DATA));
85 template <
typename TInd,
typename TData>
86 bool DoRunWithType2() {
87 auto& data = Input(DATA);
88 auto& indices = Input(INDICES);
89 auto& grad = Input(GRAD);
90 auto* output = Output(0);
92 CAFFE_ENFORCE_GE(data.ndim(), 2,
"DATA should be at least 2-D");
94 data.dim(0), grad.dim(0),
"batch sizes should be the same");
96 output->ResizeLike(data);
97 TData* out_data = output->template mutable_data<TData>();
98 if (data.size() <= 0) {
102 memset(out_data, 0, output->nbytes());
104 const TData* grad_data = grad.template data<TData>();
106 auto block_size = data.size_from_dim(2);
107 auto N = indices.size();
108 auto data_batch_size = data.size_from_dim(1);
109 auto gathered_batch_size = N * data.size_from_dim(2);
110 const TInd* idxs = indices.template data<TInd>();
112 for (
auto batch = 0; batch < grad.dim(0); ++batch) {
113 for (
auto i = 0; i < N; ++i) {
116 0 <= idx && idx < data.dim(1),
117 "INDICES element is out of DATA bounds, id=",
123 out_data + idx * block_size + batch * data_batch_size,
124 grad_data + i * block_size + batch * gathered_batch_size,
125 out_data + idx * block_size + batch * data_batch_size,
132 template <
typename TInd>
133 bool DoRunWithOtherType2() {
135 "BatchGatherGradient is not implemented on tensor of type ",
136 Input(DATA).meta().name(),
137 "Consider adding it a type in the list DispatchHelper or implementing " 138 "a generic version (which won't work for duplicated indices though)");
141 INPUT_TAGS(DATA, INDICES, GRAD);
146 #endif // CAFFE2_OPERATORS_BATCH_GATHER_OPS_H_
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...