1 #ifndef CAFFE2_OPERATORS_PACK_RNN_SEQUENCE_OP_H_ 2 #define CAFFE2_OPERATORS_PACK_RNN_SEQUENCE_OP_H_ 6 #include "caffe2/core/context.h" 7 #include "caffe2/core/operator.h" 8 #include "caffe2/utils/math.h" 12 template <
class Context,
bool Forward>
15 USE_OPERATOR_CONTEXT_FUNCTIONS;
19 bool RunOnDevice()
override {
24 template <
typename ValT>
25 bool DoRunWithType() {
28 int dim_offset = Forward ? 1 : 2;
29 auto& values = Input(0);
30 CAFFE_ENFORCE_GT(values.ndim(), dim_offset);
33 TIndex block_size = values.size_from_dim(dim_offset);
34 auto values_vec = values.template data<ValT>();
36 auto& lengths = Input(LENGTHS);
37 CAFFE_ENFORCE_EQ(lengths.ndim(), 1);
38 const auto cols = lengths.size();
39 const int32_t* lengths_vec = lengths.template data<int32_t>();
43 cols ? *std::max_element(lengths_vec, lengths_vec + cols) : 0;
44 CAFFE_ENFORCE_GE(rows, 0);
47 math::Sum<int, Context>(cols, lengths_vec, &length_sum, &context_);
54 shape.push_back(rows);
55 shape.push_back(cols);
57 shape.push_back(length_sum);
61 shape.end(), values.dims().begin() + dim_offset, values.dims().end());
63 auto* output = Output(OUTPUTVALUE);
64 output->Resize(shape);
66 auto output_data = output->template mutable_data<ValT>();
69 math::Set<ValT, Context>(output->size(), 0, output_data, &context_);
72 for (
int c = 0; c < cols; c++) {
73 for (
int r = 0; r < lengths_vec[c]; r++) {
74 auto input_offset = Forward ? (offset + r) : (r * cols + c);
75 auto output_offset = Forward ? (r * cols + c) : (offset + r);
76 context_.template CopyItems<Context, Context>(
79 values_vec + input_offset * block_size,
80 output_data + output_offset * block_size);
82 offset += lengths_vec[c];
88 INPUT_TAGS(INPUTVALUE, LENGTHS);
89 OUTPUT_TAGS(OUTPUTVALUE);
93 #endif // CAFFE2_OPERATORS_PACK_RNN_SEQUENCE_OP_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 ...