Add byte_weight_dequant_op

Summary: Pull Request resolved: pytorch#9541 Reviewed By: hlu1 Differential Revision: D8882964 fbshipit-source-id: 06d2e0d227ea6a4a8dc5ef1ea9dd1d449c149b47

Author: Jokeren

caffe2/operators/byte_weight_dequant_op.cc

@@ -0,0 +1,11 @@
+#include "caffe2/operators/byte_weight_dequant_op.h"
+
+#include "caffe2/utils/math.h"
+
+namespace caffe2 {
+
+REGISTER_CPU_OPERATOR(ByteWeightDequant, ByteWeightDequantOp<CPUContext>);
+
+OPERATOR_SCHEMA(ByteWeightDequant).NumInputs(1).NumOutputs(1);
+
+} // namespace caffe2

caffe2/operators/byte_weight_dequant_op.h

@@ -0,0 +1,55 @@
+#ifndef CAFFE2_OPERATORS_BYTE_WEIGHT_DEQUANT_OP_H_
+#define CAFFE2_OPERATORS_BYTE_WEIGHT_DEQUANT_OP_H_
+
+#include "caffe2/core/operator.h"
+#include "caffe2/utils/eigen_utils.h"
+#include "caffe2/utils/math.h"
+
+namespace caffe2 {
+
+template <typename Context>
+class ByteWeightDequantOp : public Operator<Context> {
+ public:
+ ByteWeightDequantOp(const OperatorDef& operator_def, Workspace* ws)
+ : Operator<Context>(operator_def, ws),
+ min_(OperatorBase::GetSingleArgument<float>("min", -3)),
+ max_(OperatorBase::GetSingleArgument<float>("max", 3)),
+ shape_(OperatorBase::GetRepeatedArgument<int>("shape")) {}
+
+ USE_OPERATOR_FUNCTIONS(Context);
+ using Operator<Context>::Operator;
+
+ bool RunOnDevice() override {
+ const auto& WI = Input(0);
+ auto* Y = Output(0);
+ Y->Resize(shape_);
+ float bin_interval = (max_ - min_) / 255.0;
+ int total = 1;
+ for (int i = 0; i < shape_.size(); i++) {
+ total *= Y->dim(i);
+ }
+ const uint8_t* Xdata;
+ if (WI.template IsType<uint8_t>()) {
+ CAFFE_ENFORCE(total, WI.nbytes());
+ Xdata = WI.template data<uint8_t>();
+ } else {
+ CAFFE_ENFORCE(total, WI.template data<std::string>()[0].size());
+ Xdata = reinterpret_cast<const uint8_t*>(
+ WI.template data<std::string>()[0].c_str());
+ }
+ auto* Ydata = Y->template mutable_data<float>();
+ ConstEigenVectorMap<uint8_t> index(&Xdata[0], total);
+ EigenVectorMap<float> weights(&Ydata[0], total);
+ weights = (index.cast<float>().array() * bin_interval) + min_;
+ return true;
+ }
+
+ private:
+ float min_;
+ float max_;
+ std::vector<int> shape_;
+};
+
+} // namespace caffe2
+
+#endif // CAFFE2_OPERATORS_BYTE_WEIGHT_DEQUANT_OP_H_