[Quantization] ops: 'quantize', 'dequantize'

src/operator/contrib/dequantize-inl.h

@@ -0,0 +1,99 @@
+/*!
+ *  Copyright (c) 2017 by Contributors
+ * \file dequantize-inl.h
+ * \brief Implementation of dequantize operation
+ */
+#ifndef MXNET_OPERATOR_CONTRIB_DEQUANTIZE_INL_H_
+#define MXNET_OPERATOR_CONTRIB_DEQUANTIZE_INL_H_
+
+#include <mxnet/operator_util.h>
+#include <vector>
+#include <limits>
+#include "../elemwise_op_common.h"
+#include "../mshadow_op.h"
+#include "../mxnet_op.h"
+
+namespace mxnet {
+namespace op {
+
+struct DequantizeParam : public dmlc::Parameter<DequantizeParam> {
+  int out_type;
+  DMLC_DECLARE_PARAMETER(DequantizeParam) {
+    DMLC_DECLARE_FIELD(out_type)
+    .add_enum("float32", mshadow::kFloat32)
+    .describe("Output data type.");
+  }
+};
+
+struct dequantize {
+  template<typename DstDType, typename SrcDType>
+  MSHADOW_XINLINE static void Map(int i, DstDType *out, const SrcDType *in,
+                                  float *imin_range, float *imax_range,
+                                  double imin_limit, double imax_limit,
+                                  float half_range) {
+    float scale = (*imax_range - *imin_range) / (imax_limit - imin_limit);
+    out[i] = static_cast<DstDType>((in[i] + half_range) * scale + *imin_range);
+  }
+};
+
+template<typename xpu>
+void DequantizeCompute(const nnvm::NodeAttrs& attrs,
+                     const OpContext& ctx,
+                     const std::vector<TBlob>& inputs,
+                     const std::vector<OpReqType>& req,
+                     const std::vector<TBlob>& outputs) {
+  using namespace mshadow;
+  using namespace mxnet_op;
+  Stream<xpu> *s = ctx.get_stream<xpu>();
+
+  const DequantizeParam& param = nnvm::get<DequantizeParam>(attrs.parsed);
+  MSHADOW_TYPE_SWITCH(outputs[0].type_flag_, DstDType, {
+  MSHADOW_TYPE_SWITCH(inputs[0].type_flag_, SrcDType, {
+    double min_limit = static_cast<double>(std::numeric_limits<SrcDType>::min());
+    double max_limit = static_cast<double>(std::numeric_limits<SrcDType>::max());
+    float half_range = !std::is_signed<SrcDType>::value
+      ? 0.0f
+      : (max_limit - min_limit + 1) / 2.0;
+
+    Kernel<dequantize, xpu>::Launch(s, outputs[0].Size(), outputs[0].dptr<DstDType>(),
+      inputs[0].dptr<SrcDType>(), inputs[1].dptr<float>(), inputs[2].dptr<float>(),
+      min_limit, max_limit, half_range);
+  });
+  });
+}
+
+inline bool DequantizeShape(const nnvm::NodeAttrs& attrs,
+                          std::vector<TShape> *in_attrs,
+                          std::vector<TShape> *out_attrs) {
+  const DequantizeParam& param = nnvm::get<DequantizeParam>(attrs.parsed);
+  CHECK_EQ(in_attrs->size(), 3U);
+  CHECK_EQ(out_attrs->size(), 1U);
+
+  CHECK(!shape_is_none(in_attrs->at(0)));
+  for (size_t i = 1; i < 3; ++i) {
+    CHECK(shape_is_scalar(in_attrs->at(i))) << in_attrs->at(i);
+  }
+
+  SHAPE_ASSIGN_CHECK(*out_attrs, 0, in_attrs->at(0));
+  return true;
+}
+
+inline bool DequantizeType(const nnvm::NodeAttrs& attrs,
+                         std::vector<int> *in_attrs,
+                         std::vector<int> *out_attrs) {
+  const DequantizeParam& param = nnvm::get<DequantizeParam>(attrs.parsed);
+  CHECK_EQ(in_attrs->size(), 3U);
+  CHECK_EQ(out_attrs->size(), 1U);
+  CHECK_EQ((*in_attrs)[0], mshadow::kUint8)
+    << "`dequantize` only supports uint8 input for now";
+  CHECK_EQ((*in_attrs)[1], mshadow::kFloat32)
+    << "the second input of `dequantize` should be a tensor with type of float";
+  CHECK_EQ((*in_attrs)[2], mshadow::kFloat32)
+    << "the third input of `dequantize` should be a tensor with type of float";
+  TYPE_ASSIGN_CHECK(*out_attrs, 0, mshadow::kFloat32);
+  return (*in_attrs)[0] != -1;
+}
+
+}  // namespace op
+}  // namespace mxnet
+#endif  // MXNET_OPERATOR_CONTRIB_DEQUANTIZE_INL_H_

src/operator/contrib/dequantize.cc

@@ -0,0 +1,38 @@
+/*!
+ *  Copyright (c) 2017 by Contributors
+ * \file dequantize.cc
+ * \brief
+ */
+#include "./dequantize-inl.h"
+
+namespace mxnet {
+namespace op {
+DMLC_REGISTER_PARAMETER(DequantizeParam);
+
+NNVM_REGISTER_OP(_contrib_dequantize)
+.describe(R"code(Dequantize the input tensor into a float tensor.
+[min_range, max_range] are scalar floats that spcify the range for
+the output data.
+
+Each value of the tensor will undergo the following:
+
+`out[i] = min_range + (in[i] * (max_range - min_range) / range(INPUT_TYPE))`
+
+here `range(T) = numeric_limits<T>::max() - numeric_limits<T>::min()`
+)code" ADD_FILELINE)
+.set_attr_parser(ParamParser<DequantizeParam>)
+.set_num_inputs(3)
+.set_num_outputs(1)
+.set_attr<nnvm::FInferShape>("FInferShape", DequantizeShape)
+.set_attr<nnvm::FInferType>("FInferType", DequantizeType)
+.set_attr<FCompute>("FCompute<cpu>", DequantizeCompute<cpu>)
+.set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseNone{"_dequantize"})
+.add_argument("input", "NDArray-or-Symbol", "A ndarray/symbol of type `uint8`")
+.add_argument("min_range", "NDArray-or-Symbol", "The minimum scalar value "
+  "possibly produced for the input")
+.add_argument("max_range", "NDArray-or-Symbol", "The maximum scalar value "
+  "possibly produced for the input")
+.add_arguments(DequantizeParam::__FIELDS__());
+
+}  // namespace op
+}  // namespace mxnet

src/operator/contrib/dequantize.cu

@@ -0,0 +1,15 @@
+/*!
+ *  Copyright (c) 2017 by Contributors
+ * \file dequantize.cu
+ * \brief
+ */
+#include "./dequantize-inl.h"
+
+namespace mxnet {
+namespace op {
+
+NNVM_REGISTER_OP(_contrib_dequantize)
+.set_attr<FCompute>("FCompute<gpu>", DequantizeCompute<gpu>);
+
+}  // namespace op
+}  // namespace mxnet

src/operator/contrib/quantize-inl.h

@@ -0,0 +1,101 @@
+/*!
+ *  Copyright (c) 2017 by Contributors
+ * \file quantize-inl.h
+ * \brief implementation of quantize operation
+ */
+#ifndef MXNET_OPERATOR_CONTRIB_QUANTIZE_INL_H_
+#define MXNET_OPERATOR_CONTRIB_QUANTIZE_INL_H_
+
+#include <mxnet/operator_util.h>
+#include <vector>
+#include <limits>
+#include "../elemwise_op_common.h"
+#include "../mshadow_op.h"
+#include "../mxnet_op.h"
+
+namespace mxnet {
+namespace op {
+
+struct QuantizeParam : public dmlc::Parameter<QuantizeParam> {
+  int   out_type;
+  DMLC_DECLARE_PARAMETER(QuantizeParam) {
+    DMLC_DECLARE_FIELD(out_type)
+    .add_enum("uint8", mshadow::kUint8)
+    .set_default(mshadow::kUint8)
+    .describe("Output data type.");
+  }
+};
+
+struct quantize {
+  template<typename DstDType, typename SrcDType>
+  MSHADOW_XINLINE static void Map(int i, DstDType *out, float *omin_range,
+                                  float *omax_range, const SrcDType *in,
+                                  const float *imin_range, const float *imax_range,
+                                  double min_limit, double max_limit) {
+    float scale = (max_limit - min_limit) / (*imax_range - *imin_range);
+    out[i] = static_cast<DstDType>((in[i] - *imin_range) * scale + 0.5);
+    *omin_range = *imin_range;
+    *omax_range = *imax_range;
+  }
+};
+
+template<typename xpu>
+void QuantizeCompute(const nnvm::NodeAttrs& attrs,
+                     const OpContext& ctx,
+                     const std::vector<TBlob>& inputs,
+                     const std::vector<OpReqType>& req,
+                     const std::vector<TBlob>& outputs) {
+  using namespace mshadow;
+  using namespace mxnet_op;
+  Stream<xpu> *s = ctx.get_stream<xpu>();
+
+  const QuantizeParam& param = nnvm::get<QuantizeParam>(attrs.parsed);
+  MSHADOW_TYPE_SWITCH(outputs[0].type_flag_, DstDType, {
+  MSHADOW_TYPE_SWITCH(inputs[0].type_flag_, SrcDType, {
+    Kernel<quantize, xpu>::Launch(s, outputs[0].Size(),
+      outputs[0].dptr<DstDType>(), outputs[1].dptr<float>(), outputs[2].dptr<float>(),
+      inputs[0].dptr<SrcDType>(), inputs[1].dptr<float>(), inputs[2].dptr<float>(),
+      std::numeric_limits<DstDType>::min(), std::numeric_limits<DstDType>::max());
+  });
+  });
+}
+
+inline bool QuantizeShape(const nnvm::NodeAttrs& attrs,
+                          std::vector<TShape> *in_attrs,
+                          std::vector<TShape> *out_attrs) {
+  const QuantizeParam& param = nnvm::get<QuantizeParam>(attrs.parsed);
+  CHECK_EQ(in_attrs->size(), 3U);
+  CHECK_EQ(out_attrs->size(), 3U);
+
+  CHECK(!shape_is_none(in_attrs->at(0)));
+  for (size_t i = 1; i < 3; ++i) {
+    CHECK(shape_is_scalar(in_attrs->at(i)));
+  }
+
+  SHAPE_ASSIGN_CHECK(*out_attrs, 0, in_attrs->at(0));
+  SHAPE_ASSIGN_CHECK(*out_attrs, 1, TShape{1});
+  SHAPE_ASSIGN_CHECK(*out_attrs, 2, TShape{1});
+  return true;
+}
+
+inline bool QuantizeType(const nnvm::NodeAttrs& attrs,
+                         std::vector<int> *in_attrs,
+                         std::vector<int> *out_attrs) {
+  const QuantizeParam& param = nnvm::get<QuantizeParam>(attrs.parsed);
+  CHECK_EQ(in_attrs->size(), 3U);
+  CHECK_EQ(out_attrs->size(), 3U);
+  CHECK_EQ((*in_attrs)[0], mshadow::kFloat32)
+    << "`quantize` only supports float32 input for now";
+  CHECK_EQ((*in_attrs)[1], mshadow::kFloat32)
+    << "the second input of `quantize` should be a tensor with type of float";
+  CHECK_EQ((*in_attrs)[2], mshadow::kFloat32)
+    << "the third input of `quantize` should be a tensor with type of float";
+  TYPE_ASSIGN_CHECK(*out_attrs, 0, mshadow::kUint8);
+  TYPE_ASSIGN_CHECK(*out_attrs, 1, mshadow::kFloat32);
+  TYPE_ASSIGN_CHECK(*out_attrs, 2, mshadow::kFloat32);
+  return (*in_attrs)[0] != -1;
+}
+
+}  // namespace op
+}  // namespace mxnet
+#endif  // MXNET_OPERATOR_CONTRIB_QUANTIZE_INL_H_

src/operator/contrib/quantize.cc

@@ -0,0 +1,38 @@
+/*!
+ *  Copyright (c) 2017 by Contributors
+ * \file quantize.cc
+ * \brief
+ */
+#include "./quantize-inl.h"
+
+namespace mxnet {
+namespace op {
+DMLC_REGISTER_PARAMETER(QuantizeParam);
+
+NNVM_REGISTER_OP(_contrib_quantize)
+.describe(R"code(Quantize a input tensor from float to `out_type`,
+with user-specified `min_range` and `max_range`.
+
+[min_range, max_range] are scalar floats that spcify the range for
+the input data. Each value of the tensor will undergo the following:
+
+`out[i] = (in[i] - min_range) * range(OUTPUT_TYPE) / (max_range - min_range)`
+
+here `range(T) = numeric_limits<T>::max() - numeric_limits<T>::min()`
+)code" ADD_FILELINE)
+.set_attr_parser(ParamParser<QuantizeParam>)
+.set_num_inputs(3)
+.set_num_outputs(3)
+.set_attr<nnvm::FInferShape>("FInferShape", QuantizeShape)
+.set_attr<nnvm::FInferType>("FInferType", QuantizeType)
+.set_attr<FCompute>("FCompute<cpu>", QuantizeCompute<cpu>)
+.set_attr<nnvm::FGradient>("FGradient", ElemwiseGradUseNone{"_quantize"})
+.add_argument("input", "NDArray-or-Symbol", "A ndarray/symbol of type `float32`")
+.add_argument("min_range", "NDArray-or-Symbol", "The minimum scalar value "
+  "possibly produced for the input")
+.add_argument("max_range", "NDArray-or-Symbol", "The maximum scalar value "
+  "possibly produced for the input")
+.add_arguments(QuantizeParam::__FIELDS__());
+
+}  // namespace op
+}  // namespace mxnet

src/operator/contrib/quantize.cu

@@ -0,0 +1,15 @@
+/*!
+ *  Copyright (c) 2017 by Contributors
+ * \file quantize.cu
+ * \brief
+ */
+#include "./quantize-inl.h"
+
+namespace mxnet {
+namespace op {
+
+NNVM_REGISTER_OP(_contrib_quantize)
+.set_attr<FCompute>("FCompute<gpu>", QuantizeCompute<gpu>);
+
+}  // namespace op
+}  // namespace mxnet

src/operator/operator_common.h

@@ -77,6 +77,11 @@ inline bool type_is_none(const int& x) {
   return x == -1;
 }
 
+/*! \brief check if shape is scalar({1}). */
+inline bool shape_is_scalar(const TShape& x) {
+  return x.ndim() == 1 && x.Size() == 1;
+}
+
 /*! \brief get string representation of shape */
 inline std::string shape_string(const TShape& x) {
   std::ostringstream os;

tests/python/unittest/test_operator.py

@@ -2991,6 +2991,18 @@ def test_pick_helper(index_type=np.int32):
     test_pick_helper(np.int32)
     test_pick_helper(np.float32)
 
+def test_quantization_op():
+  min0 = mx.nd.array([0.0])
+  max0 = mx.nd.array([1.0])
+  a  = mx.nd.array([[0.1392, 0.5928], [0.6027, 0.8579]])
+  qa, min1, max1 = mx.contrib.nd.quantize(a, min0, max0, out_type='uint8')
+  a_ = mx.contrib.nd.dequantize(qa, min1, max1, out_type='float32')
+
+  qa_real = mx.nd.array([[35, 151], [154, 219]])
+  a_real  = mx.nd.array([[0.13725491, 0.59215689], [0.60392159, 0.8588236]])
+
+  assert same(qa.asnumpy(), qa_real.asnumpy())
+  assert same(a_.asnumpy(),  a_real.asnumpy())
 
 def test_custom_op():
     class Sqr(mx.operator.CustomOp):
@@ -3098,5 +3110,6 @@ def create_operator(self, ctx, shapes, dtypes):
     test_tile()
     test_one_hot()
     test_where()
+    test_quantization_op()
     test_relu()
     test_sigmoid()