1 #ifndef CAFFE2_OPERATORS_REDUCTION_FRONT_BACK_OPS_H_ 2 #define CAFFE2_OPERATORS_REDUCTION_FRONT_BACK_OPS_H_ 4 #include "caffe2/core/context.h" 5 #include "caffe2/core/logging.h" 6 #include "caffe2/core/operator.h" 10 template <
class Context,
bool FIRSTDIMS,
bool NORMALIZE>
16 OperatorBase::GetSingleArgument<int32_t>(
"num_reduce_dim", 1)) {}
18 USE_OPERATOR_CONTEXT_FUNCTIONS;
20 bool RunOnDevice()
override {
26 bool DoRunWithType() {
31 num_reduce_dims_ >= 0 && num_reduce_dims_ <= X.dims().size(),
32 "For N-dim input tensor, support num_reduce_dims in range [0, N].");
34 vector<TIndex> output_shape;
35 int start_index = FIRSTDIMS ? num_reduce_dims_ : 0;
37 FIRSTDIMS ? X.dims().size() : X.dims().size() - num_reduce_dims_;
38 for (
int i = start_index; i < end_index; ++i) {
39 output_shape.push_back(X.dims()[i]);
41 Y->Resize(output_shape);
43 const int rows = FIRSTDIMS ? X.size_to_dim(num_reduce_dims_)
44 : X.size_to_dim(X.ndim() - num_reduce_dims_);
45 const int cols = FIRSTDIMS ? X.size_from_dim(num_reduce_dims_)
46 : X.size_from_dim(X.ndim() - num_reduce_dims_);
48 if (cols == 0 || rows == 0) {
52 const int32_t* lengths_data =
nullptr;
53 if (InputSize() > 1) {
54 const auto& lengths = Input(1);
55 lengths_data = lengths.template data<int32_t>();
57 num_reduce_dims_ == 1,
58 "Given lengths input, the number of reduce dimensions should be one.");
59 const int batch_size = FIRSTDIMS ? cols : rows;
61 lengths.size() == batch_size,
62 "The size of lengths vector doesn't match the batch size.");
65 const T* in_data = X.template data<T>();
66 T* out_data = Y->template mutable_data<T>();
67 Compute(rows, cols, in_data, lengths_data, out_data);
78 const int32_t* lengths_data,
84 template <
class Context,
bool FIRSTDIMS,
bool NORMALIZE>
90 OperatorBase::GetSingleArgument<int32_t>(
"num_reduce_dim", 1)) {}
92 USE_OPERATOR_CONTEXT_FUNCTIONS;
94 bool RunOnDevice()
override {
100 bool DoRunWithType() {
102 auto& input_1 = Input(1);
103 auto* dX = Output(0);
109 if (input_1.ndim() == 1 && input_1.template IsType<TIndex>()) {
111 shape_.CopyFrom(input_1);
113 vector<TIndex> output_shape(
114 shape_.template data<TIndex>(),
115 shape_.template data<TIndex>() + shape_.size());
116 dX->Resize(output_shape);
119 dX->ResizeLike(input_1);
122 const int rows = FIRSTDIMS ? dX->size_to_dim(num_reduce_dims_)
123 : dX->size_to_dim(dX->ndim() - num_reduce_dims_);
124 const int cols = FIRSTDIMS
125 ? dX->size_from_dim(num_reduce_dims_)
126 : dX->size_from_dim(dX->ndim() - num_reduce_dims_);
128 const int32_t* lengths_data =
nullptr;
129 if (InputSize() > 2) {
130 const auto& lengths = Input(2);
131 lengths_data = lengths.template data<int32_t>();
133 num_reduce_dims_ == 1,
134 "Given lengths input, the number of reduce dimensions should be one.");
135 const int batch_size = FIRSTDIMS ? cols : rows;
137 lengths.size() == batch_size,
138 "The size of lengths vector doesn't match the batch size.");
141 const T* dYdata = dY.template data<T>();
142 T* dXdata = dX->template mutable_data<T>();
143 Compute<T>(rows, cols, dYdata, lengths_data, dXdata);
148 template <
typename T>
153 const int32_t* lengths_data,
155 int num_reduce_dims_;
160 template <
typename T,
class Context,
bool FIRSTDIMS>
166 OperatorBase::GetSingleArgument<int32_t>(
"num_reduce_dim", 1)) {}
168 USE_OPERATOR_CONTEXT_FUNCTIONS;
175 num_reduce_dims_ >= 0 && num_reduce_dims_ <= X.dims().size(),
176 "For N-dim input tensor, support num_reduce_dims in range [0, N].");
178 const int rows = FIRSTDIMS ? X.size_to_dim(num_reduce_dims_)
179 : X.size_to_dim(X.ndim() - num_reduce_dims_);
180 const int cols = FIRSTDIMS ? X.size_from_dim(num_reduce_dims_)
181 : X.size_from_dim(X.ndim() - num_reduce_dims_);
183 vector<TIndex> output_shape;
184 int start_index = FIRSTDIMS ? num_reduce_dims_ : 0;
186 FIRSTDIMS ? X.dims().size() : X.dims().size() - num_reduce_dims_;
188 for (
int i = start_index; i < end_index; ++i) {
189 output_shape.push_back(X.dims()[i]);
191 Y->Resize(output_shape);
193 if (cols == 0 || rows == 0) {
197 const int32_t* lengths_data =
nullptr;
198 if (InputSize() > 1) {
199 const auto& lengths = Input(1);
200 lengths_data = lengths.template data<int32_t>();
202 num_reduce_dims_ == 1,
203 "Given lengths input, the number of reduce dimensions should be one.");
204 const int batch_size = FIRSTDIMS ? cols : rows;
206 lengths.size() == batch_size,
207 "The size of lengths vector doesn't match the batch size.");
210 const float* data = X.template data<float>();
211 float* out_data = Y->template mutable_data<float>();
212 Compute(rows, cols, data, lengths_data, out_data);
221 const int32_t* lengths_data,
224 int num_reduce_dims_;
227 template <
typename T,
class Context,
bool FIRSTDIMS>
233 OperatorBase::GetSingleArgument<int32_t>(
"num_reduce_dim", 1)) {}
235 USE_OPERATOR_CONTEXT_FUNCTIONS;
237 bool RunOnDevice()
override {
241 auto* dX = Output(0);
244 const int rows = FIRSTDIMS ? X.size_to_dim(num_reduce_dims_)
245 : X.size_to_dim(X.ndim() - num_reduce_dims_);
246 const int cols = FIRSTDIMS ? X.size_from_dim(num_reduce_dims_)
247 : X.size_from_dim(X.ndim() - num_reduce_dims_);
249 const float* dYdata = dY.template data<float>();
250 const float* Xdata = X.template data<float>();
251 const float* Ydata = Y.template data<float>();
253 const int32_t* lengths_data =
nullptr;
254 if (InputSize() > 3) {
255 const auto& lengths = Input(3);
256 lengths_data = lengths.template data<int32_t>();
258 num_reduce_dims_ == 1,
259 "Given lengths input, the number of reduce dimensions should be one.");
260 const int batch_size = FIRSTDIMS ? cols : rows;
262 lengths.size() == batch_size,
263 "The size of lengths vector doesn't match the batch size.");
266 float* dXdata = dX->template mutable_data<float>();
267 Compute(rows, cols, dYdata, Xdata, Ydata, lengths_data, dXdata);
278 const int32_t* lengths_data,
281 int num_reduce_dims_;
286 #endif // CAFFE2_OPERATORS_REDUCTION_FRONT_BACK_OPS_H_
Workspace is a class that holds all the related objects created during runtime: (1) all blobs...
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...