Caffe2 - C++ API
A deep learning, cross platform ML framework
int8_flatten_op.h
1 #ifndef CAFFE2_OPERATORS_INT8_FLATTEN_OP_H_
2 #define CAFFE2_OPERATORS_INT8_FLATTEN_OP_H_
3 
4 #include "caffe2/core/context.h"
5 #include "caffe2/core/operator.h"
6 #include "caffe2/core/tensor_int8.h"
7 #include "caffe2/operators/quantized/int8_utils.h"
8 
9 namespace caffe2 {
10 
11 namespace int8 {
12 
13 class Int8FlattenOp : public Operator<CPUContext> {
14  public:
15  template <class... Args>
16  explicit Int8FlattenOp(Args&&... args)
17  : Operator<CPUContext>(std::forward<Args>(args)...),
18  axis_(this->template GetSingleArgument<int>("axis", 1)) {}
19 
20  bool RunOnDevice() override {
21  const auto& X = Inputs()[0]->Get<Int8TensorCPU>();
22  auto* Y = Outputs()[0]->GetMutable<Int8TensorCPU>();
23  int32_t Y_offset = this->template GetSingleArgument<int>("Y_zero_point", 0);
24  auto Y_scale = this->template GetSingleArgument<float>("Y_scale", 1);
25  CHECK_EQ(Y_offset, X.zero_point);
26  CHECK_EQ(Y_scale, X.scale);
27  Y->scale = Y_scale;
28  Y->zero_point = Y_offset;
29  CAFFE_ENFORCE_GE(
30  X.t.sizes().size(), axis_, "The rank of the tensor must be >= axis.");
31  Y->t.Resize(X.t.size_to_dim(axis_), X.t.size_from_dim(axis_));
32  context_.CopyItemsToCPU(
33  X.t.dtype(),
34  X.t.numel(),
35  X.t.raw_data(),
36  Y->t.raw_mutable_data(X.t.dtype()));
37  return true;
38  }
39 
40  private:
41  int axis_;
42 };
43 
44 } // namespace int8
45 
46 } // namespace caffe2
47 
48 #endif // CAFFE2_OPERATORS_INT8_FLATTEN_OP_H_
A global dictionary that holds information about what Caffe2 modules have been loaded in the current ...
Definition: blob.h:13