3 #include "caffe2/core/operator.h" 4 #include "caffe2/utils/math.h" 8 template <
class Context>
11 USE_OPERATOR_CONTEXT_FUNCTIONS;
14 bool RunOnDevice()
override {
16 this, OperatorBase::Input<TensorCPU>(INDICES));
19 template <
typename Index>
20 bool DoRunWithType() {
21 const auto& data = Input(DATA);
22 const auto& indices = Input(INDICES);
23 auto* output = Output(0);
25 CAFFE_ENFORCE_EQ(data.ndim(), 2,
"DATA must be a matrix");
26 CAFFE_ENFORCE_EQ(indices.ndim(), 1,
"INDICES must be a vector");
27 CAFFE_ENFORCE_GT(data.dim(1), 8,
"DATA must have more than 8 columns");
30 const std::vector<TIndex> shape = {indices.dim(0), data.dim(1) - 8};
31 output->Resize(shape);
33 int block_size = shape[1];
34 auto block_bytesize = data.size_from_dim(1) * data.meta().itemsize();
35 int N = indices.size();
37 const uint8_t* src_base = data.template data<uint8_t>();
38 const Index* idxs = indices.template data<Index>();
39 auto out = output->template mutable_data<float>();
41 for (
int i = 0; i < N; ++i) {
44 0 <= idx && idx < data.dim(0),
45 "INDICES element is out of DATA bounds, id=",
49 const uint8_t* src = src_base + idx * block_bytesize;
50 ConstEigenVectorArrayMap<uint8_t> input_row_values(src, shape[1]);
51 ConstEigenVectorArrayMap<float> input_row_scale_bias(
52 reinterpret_cast<const float*>(src + shape[1]), 2);
54 EigenVectorArrayMap<float> output_row(out + i * shape[1], shape[1]);
56 output_row = input_row_values.cast<
float>() * input_row_scale_bias(0) +
57 input_row_scale_bias(1);
62 INPUT_TAGS(DATA, INDICES);
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...