Caffe2 - C++ API
A deep learning, cross platform ML framework
batch_gather_ops.h
1 
17 #ifndef CAFFE2_OPERATORS_BATCH_GATHER_OPS_H_
18 #define CAFFE2_OPERATORS_BATCH_GATHER_OPS_H_
19 
20 #include "caffe2/core/context.h"
21 #include "caffe2/core/operator.h"
22 #include "caffe2/utils/math.h"
23 
24 namespace caffe2 {
25 
26 template <class Context>
27 class BatchGatherOp final : public Operator<Context> {
28  public:
29  USE_OPERATOR_CONTEXT_FUNCTIONS;
30  USE_SIMPLE_CTOR_DTOR(BatchGatherOp)
31 
32  bool RunOnDevice() override {
34  this, OperatorBase::Input<TensorCPU>(INDICES));
35  }
36 
37  template <typename TInd>
38  bool DoRunWithType() {
39  auto& data = Input(DATA);
40  auto& indices = Input(INDICES);
41  auto* output = Output(0);
42 
43  CAFFE_ENFORCE_GE(data.ndim(), 2, "DATA should be at least 2-D");
44 
45  vector<TIndex> shape;
46  shape.push_back(data.dim(0));
47  shape.insert(shape.end(), indices.dims().begin(), indices.dims().end());
48  shape.insert(shape.end(), data.dims().begin() + 2, data.dims().end());
49  output->Resize(shape);
50 
51  auto block_size = data.size_from_dim(2);
52  auto block_bytesize = block_size * data.meta().itemsize();
53  auto N = indices.size();
54  auto data_batch_bytesize = data.size_from_dim(1) * data.meta().itemsize();
55  auto gathered_batch_bytesize =
56  N * data.size_from_dim(2) * data.meta().itemsize();
57  const TInd* idxs = indices.template data<TInd>();
58  auto src_base = static_cast<const char*>(data.raw_data());
59  auto out = static_cast<char*>(output->raw_mutable_data(data.meta()));
60 
61  for (auto batch = 0; batch < data.dim(0); ++batch) {
62  for (auto i = 0; i < N; ++i) {
63  auto idx = idxs[i];
64  CAFFE_ENFORCE(
65  0 <= idx && idx < data.dim(1),
66  "INDICES element is out of DATA bounds, id=",
67  idx,
68  " data_dim=",
69  data.dim(1));
70  auto src =
71  src_base + idx * block_bytesize + batch * data_batch_bytesize;
72  auto dst = out + i * block_bytesize + batch * gathered_batch_bytesize;
73  context_.template CopyItems<Context, Context>(
74  data.meta(), block_size, src, dst);
75  }
76  }
77  return true;
78  }
79 
80  INPUT_TAGS(DATA, INDICES);
81 };
82 
83 template <class Context>
84 class BatchGatherGradientOp final : public Operator<Context> {
85  public:
86  USE_OPERATOR_CONTEXT_FUNCTIONS;
87  USE_SIMPLE_CTOR_DTOR(BatchGatherGradientOp);
88 
89  bool RunOnDevice() override {
91  this, OperatorBase::Input<TensorCPU>(INDICES));
92  }
93 
94  template <typename TInd>
95  bool DoRunWithType() {
96  return DispatchHelper<
98  TInd>::call(this, Input(DATA));
99  }
100 
101  template <typename TInd, typename TData>
102  bool DoRunWithType2() {
103  auto& data = Input(DATA);
104  auto& indices = Input(INDICES);
105  auto& grad = Input(GRAD);
106  auto* output = Output(0);
107 
108  CAFFE_ENFORCE_GE(data.ndim(), 2, "DATA should be at least 2-D");
109  CAFFE_ENFORCE_EQ(
110  data.dim(0), grad.dim(0), "batch sizes should be the same");
111 
112  output->ResizeLike(data);
113  TData* out_data = output->template mutable_data<TData>();
114  memset(out_data, 0, output->nbytes());
115 
116  const TData* grad_data = grad.template data<TData>();
117 
118  auto block_size = data.size_from_dim(2);
119  auto N = indices.size();
120  auto data_batch_size = data.size_from_dim(1);
121  auto gathered_batch_size = N * data.size_from_dim(2);
122  const TInd* idxs = indices.template data<TInd>();
123 
124  for (auto batch = 0; batch < grad.dim(0); ++batch) {
125  for (auto i = 0; i < N; ++i) {
126  auto idx = idxs[i];
127  CAFFE_ENFORCE(
128  0 <= idx && idx < data.dim(1),
129  "INDICES element is out of DATA bounds, id=",
130  idx,
131  " data_dim=",
132  data.dim(1));
133  math::Add(
134  block_size,
135  out_data + idx * block_size + batch * data_batch_size,
136  grad_data + i * block_size + batch * gathered_batch_size,
137  out_data + idx * block_size + batch * data_batch_size,
138  &context_);
139  }
140  }
141  return true;
142  }
143 
144  template <typename TInd>
145  bool DoRunWithOtherType2() {
146  CAFFE_THROW(
147  "BatchGatherGradient is not implemented on tensor of type ",
148  Input(DATA).meta().name(),
149  "Consider adding it a type in the list DispatchHelper or implementing "
150  "a generic version (which won't work for duplicated indices though)");
151  }
152 
153  INPUT_TAGS(DATA, INDICES, GRAD);
154 };
155 
156 } // namespace caffe2
157 
158 #endif // CAFFE2_OPERATORS_BATCH_GATHER_OPS_H_
Copyright (c) 2016-present, Facebook, Inc.