File: int8_quantize_op.cc

package info (click to toggle)
pytorch 1.7.1-7
  • links: PTS, VCS
  • area: main
  • in suites: bullseye
  • size: 80,340 kB
  • sloc: cpp: 670,830; python: 343,991; ansic: 67,845; asm: 5,503; sh: 2,924; java: 2,888; xml: 266; makefile: 244; ruby: 148; yacc: 144; objc: 51; lex: 44
file content (60 lines) | stat: -rw-r--r-- 1,644 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
#include <caffe2/ideep/ideep_utils.h>

using namespace caffe2;

namespace {

class IDEEPInt8QuantizeOp final : public IDEEPOperator {
 public:
  USE_IDEEP_DEF_ALIASES();
  USE_IDEEP_OPERATOR_FUNCTIONS();

  IDEEPInt8QuantizeOp(const OperatorDef& operator_def, Workspace* ws)
      : IDEEPOperator(operator_def, ws),
        scale_(this->template GetSingleArgument<float>("Y_scale", 1.0)),
        zero_point_(
            this->template GetSingleArgument<int32_t>("Y_zero_point", 0)) {

    if (HasArgument("output_order")) {
      Y_fmt_ = static_cast<iformat>(
        this->template GetSingleArgument<int>("output_order",
                                              static_cast<int>(iformat::nchw)));
    }

    CAFFE_ENFORCE(zero_point_ == 0 || zero_point_ == 128,
        "Not support this zero point");
    Y_data_type_ = zero_point_ == 0 ? idtype::u8 : idtype::s8;
    Y_scales_ = ConvertScales({scale_});
  }
  virtual ~IDEEPInt8QuantizeOp() {}

  bool RunOnDevice() override {
    const auto& X = Input(0);
    CAFFE_ENFORCE(X.get_data_type() == idtype::f32, "Not support data type");

    auto* Y = Output(0);
    if (Y_fmt_ != iformat::undef) {
      Y->init(X.get_desc().to_type(Y_data_type_).to_format(Y_fmt_));
    } else {
      Y->init(X.get_desc().to_type(Y_data_type_));
    }
    Y->set_scale(Y_scales_);
    Y->feed_from(X);

    return true;
  }

 private:
  float scale_;
  int32_t zero_point_;
  iscale Y_scales_;
  idtype Y_data_type_;
  iformat Y_fmt_ {iformat::undef};

  INPUT_TAGS(INPUT0);
  OUTPUT_TAGS(OUTPUT);
};

REGISTER_IDEEP_OPERATOR_WITH_ENGINE(Int8Quantize, DNNLOWP, IDEEPInt8QuantizeOp);

} // namespace