marcoabreu closed pull request #10433: [MXNET-290] MKLDNN support for model
quantization
URL: https://github.com/apache/incubator-mxnet/pull/10433
This is a PR merged from a forked repository.
As GitHub hides the original diff on merge, it is displayed below for
the sake of provenance:
As this is a foreign pull request (from a fork), the diff is supplied
below (as it won't show otherwise due to GitHub magic):
diff --git a/ci/docker/runtime_functions.sh b/ci/docker/runtime_functions.sh
index 7abe767c869..b2f0d8ddfcb 100755
--- a/ci/docker/runtime_functions.sh
+++ b/ci/docker/runtime_functions.sh
@@ -381,13 +381,12 @@ unittest_ubuntu_python3_cpu() {
unittest_ubuntu_python3_cpu_mkldnn() {
set -ex
- export PYTHONPATH=./python/
+ export PYTHONPATH=./python/
# MXNET_MKLDNN_DEBUG is buggy and produces false positives
# https://github.com/apache/incubator-mxnet/issues/10026
#export MXNET_MKLDNN_DEBUG=1 # Ignored if not present
export MXNET_STORAGE_FALLBACK_LOG_VERBOSE=0
nosetests-3.4 --verbose tests/python/unittest
- nosetests-3.4 --verbose tests/python/quantization
nosetests-3.4 --verbose tests/python/mkl
}
diff --git a/example/quantization/imagenet_gen_qsym.py
b/example/quantization/imagenet_gen_qsym.py
index 045ce62489a..85474b663fa 100644
--- a/example/quantization/imagenet_gen_qsym.py
+++ b/example/quantization/imagenet_gen_qsym.py
@@ -53,6 +53,7 @@ def save_params(fname, arg_params, aux_params, logger=None):
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Generate a calibrated
quantized model from a FP32 model')
+ parser.add_argument('--ctx', type=str, default='gpu')
parser.add_argument('--model', type=str, choices=['imagenet1k-resnet-152',
'imagenet1k-inception-bn'],
help='currently only supports imagenet1k-resnet-152 or
imagenet1k-inception-bn')
parser.add_argument('--batch-size', type=int, default=32)
@@ -91,8 +92,18 @@ def save_params(fname, arg_params, aux_params, logger=None):
' thresholds. This mode is expected to produce
the best inference accuracy of all three'
' kinds of quantized models if the calibration
dataset is representative enough of the'
' inference dataset.')
+ parser.add_argument('--quantized-dtype', type=str, default='int8',
+ choices=['int8', 'uint8'],
+ help='quantization destination data type for input
data')
args = parser.parse_args()
+ if args.ctx == 'gpu':
+ ctx = mx.gpu(0)
+ elif args.ctx == 'cpu':
+ ctx = mx.cpu(0)
+ else:
+ raise ValueError('ctx %s is not supported in this script' % args.ctx)
+
logging.basicConfig()
logger = logging.getLogger('logger')
logger.setLevel(logging.INFO)
@@ -129,17 +140,26 @@ def save_params(fname, arg_params, aux_params,
logger=None):
excluded_sym_names = []
if args.model == 'imagenet1k-resnet-152':
rgb_mean = '0,0,0'
- calib_layer = lambda name: name.endswith('_output') and
(name.find('conv') != -1
- or
name.find('sc') != -1
- or
name.find('fc') != -1)
+ if args.ctx == 'gpu':
+ calib_layer = lambda name: name.endswith('_output') and
(name.find('conv') != -1
+ or
name.find('sc') != -1
+ or
name.find('fc') != -1)
+ else:
+ calib_layer = lambda name: name.endswith('_output') and
(name.find('conv') != -1
+ or
name.find('sc') != -1)
+ excluded_sym_names += ['flatten0', 'fc1']
if exclude_first_conv:
- excluded_sym_names = ['conv0']
+ excluded_sym_names += ['conv0']
elif args.model == 'imagenet1k-inception-bn':
rgb_mean = '123.68,116.779,103.939'
- calib_layer = lambda name: name.endswith('_output') and
(name.find('conv') != -1
- or
name.find('fc') != -1)
+ if args.ctx == 'gpu':
+ calib_layer = lambda name: name.endswith('_output') and
(name.find('conv') != -1
+ or
name.find('fc') != -1)
+ else:
+ calib_layer = lambda name: name.endswith('_output') and
(name.find('conv') != -1)
+ excluded_sym_names += ['flatten', 'fc1']
if exclude_first_conv:
- excluded_sym_names = ['conv_1']
+ excluded_sym_names += ['conv_1']
else:
raise ValueError('model %s is not supported in this script' %
args.model)
@@ -156,8 +176,9 @@ def save_params(fname, arg_params, aux_params, logger=None):
if calib_mode == 'none':
logger.info('Quantizing FP32 model %s' % args.model)
qsym, qarg_params, aux_params = quantize_model(sym=sym,
arg_params=arg_params, aux_params=aux_params,
-
excluded_sym_names=excluded_sym_names,
- calib_mode=calib_mode,
logger=logger)
+ ctx=ctx,
excluded_sym_names=excluded_sym_names,
+ calib_mode=calib_mode,
quantized_dtype=args.quantized_dtype,
+ logger=logger)
sym_name = '%s-symbol.json' % (prefix + '-quantized')
save_symbol(sym_name, qsym, logger)
else:
@@ -176,10 +197,11 @@ def save_params(fname, arg_params, aux_params,
logger=None):
**mean_args)
cqsym, qarg_params, aux_params = quantize_model(sym=sym,
arg_params=arg_params, aux_params=aux_params,
- ctx=mx.gpu(0),
excluded_sym_names=excluded_sym_names,
+ ctx=ctx,
excluded_sym_names=excluded_sym_names,
calib_mode=calib_mode,
calib_data=data,
num_calib_examples=num_calib_batches * batch_size,
-
calib_layer=calib_layer, logger=logger)
+
calib_layer=calib_layer, quantized_dtype=args.quantized_dtype,
+ logger=logger)
if calib_mode == 'entropy':
suffix = '-quantized-%dbatches-entropy' % num_calib_batches
elif calib_mode == 'naive':
diff --git a/example/quantization/imagenet_inference.py
b/example/quantization/imagenet_inference.py
index fe3f2661c65..85649530aa0 100644
--- a/example/quantization/imagenet_inference.py
+++ b/example/quantization/imagenet_inference.py
@@ -99,6 +99,7 @@ def score(sym, arg_params, aux_params, data, devs,
label_name, max_num_examples,
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Score a model on a dataset')
+ parser.add_argument('--ctx', type=str, default='gpu')
parser.add_argument('--symbol-file', type=str, required=True, help='symbol
file path')
parser.add_argument('--param-file', type=str, required=True, help='param
file path')
parser.add_argument('--batch-size', type=int, default=32)
@@ -122,6 +123,13 @@ def score(sym, arg_params, aux_params, data, devs,
label_name, max_num_examples,
args = parser.parse_args()
+ if args.ctx == 'gpu':
+ ctx = mx.gpu(0)
+ elif args.ctx == 'cpu':
+ ctx = mx.cpu(0)
+ else:
+ raise ValueError('ctx %s is not supported in this script' % args.ctx)
+
logging.basicConfig()
logger = logging.getLogger('logger')
logger.setLevel(logging.INFO)
@@ -172,5 +180,5 @@ def score(sym, arg_params, aux_params, data, devs,
label_name, max_num_examples,
num_inference_images = args.num_inference_batches * batch_size
logger.info('Running model %s for inference' % symbol_file)
- score(sym, arg_params, aux_params, data, [mx.gpu(0)], label_name,
+ score(sym, arg_params, aux_params, data, [ctx], label_name,
max_num_examples=num_inference_images, logger=logger)
diff --git a/include/mxnet/c_api.h b/include/mxnet/c_api.h
index 06e39bfeb38..c01ae4943e2 100644
--- a/include/mxnet/c_api.h
+++ b/include/mxnet/c_api.h
@@ -1436,13 +1436,15 @@ MXNET_DLL int MXSymbolInferType(SymbolHandle sym,
* \param excluded_symbols array of symbols to be excluded from being quantized
* \param num_offline number of parameters that are quantized offline
* \param offline_params array of c strings representing the names of params
quantized offline
+ * \param quantized_dtype the quantized destination type for input data.
*/
MXNET_DLL int MXQuantizeSymbol(SymbolHandle sym_handle,
SymbolHandle *ret_sym_handle,
const mx_uint num_excluded_symbols,
const SymbolHandle *excluded_symbols,
const mx_uint num_offline,
- const char **offline_params);
+ const char **offline_params,
+ const char *quantized_dtype);
/*!
* \brief Set calibration table to node attributes in the sym
diff --git a/python/mxnet/contrib/quantization.py
b/python/mxnet/contrib/quantization.py
index c9c58a9c9ba..1314b97028a 100644
--- a/python/mxnet/contrib/quantization.py
+++ b/python/mxnet/contrib/quantization.py
@@ -72,7 +72,8 @@ def _quantize_params(qsym, params):
return quantized_params
-def _quantize_symbol(sym, excluded_symbols=None, offline_params=None):
+def _quantize_symbol(sym, excluded_symbols=None, offline_params=None,
+ quantized_dtype='int8'):
"""Given a symbol object representing a neural network of data type FP32,
quantize it into a INT8 network.
@@ -86,6 +87,8 @@ def _quantize_symbol(sym, excluded_symbols=None,
offline_params=None):
Names of the parameters that users want to quantize offline. It's
always recommended to
quantize parameters offline so that quantizing parameters during the
inference can be
avoided.
+ quantized_dtype: str
+ The quantized destination type for input data.
"""
num_excluded_symbols = 0
excluded_handles = []
@@ -108,7 +111,8 @@ def _quantize_symbol(sym, excluded_symbols=None,
offline_params=None):
mx_uint(num_excluded_symbols),
c_array(SymbolHandle, excluded_handles),
mx_uint(num_offline),
- c_array(ctypes.c_char_p, offline)))
+ c_array(ctypes.c_char_p, offline),
+ c_str(quantized_dtype)))
return Symbol(out)
@@ -401,7 +405,8 @@ def _load_params(params, logger=logging):
def quantize_model(sym, arg_params, aux_params,
data_names=('data',), label_names=('softmax_label',),
ctx=cpu(), excluded_sym_names=None, calib_mode='entropy',
- calib_data=None, num_calib_examples=None, calib_layer=None,
logger=logging):
+ calib_data=None, num_calib_examples=None, calib_layer=None,
+ quantized_dtype='int8', logger=logging):
"""User-level API for generating a quantized model from a FP32 model w/ or
w/o calibration.
The backend quantized operators are only enabled for Linux systems. Please
do not run
inference using the quantized models on Windows for now.
@@ -451,6 +456,9 @@ def quantize_model(sym, arg_params, aux_params,
calibrate this layer. If yes, the statistics of the layer's output
will be collected;
otherwise, no information of the layer's output will be collected. If
not provided,
all the layers' outputs that need requantization will be collected.
+ quantized_dtype : str
+ The quantized destination type for input data. Currently support 'int8'
+ and 'uint8', default value is 'int8'.
logger : Object
A logging object for printing information during the process of
quantization.
@@ -473,8 +481,13 @@ def quantize_model(sym, arg_params, aux_params,
idx = nodes.list_outputs().index(sym_name + '_output')
excluded_syms.append(nodes[idx])
logger.info('Quantizing symbol')
+
+ if quantized_dtype != 'int8' and quantized_dtype != 'uint8':
+ raise ValueError('unknown quantized_dtype %s received,'
+ ' expected `int8` or `uint8`' % quantized_dtype)
qsym = _quantize_symbol(sym, excluded_symbols=excluded_syms,
- offline_params=list(arg_params.keys()))
+ offline_params=list(arg_params.keys()),
+ quantized_dtype=quantized_dtype)
logger.info('Quantizing parameters')
qarg_params = _quantize_params(qsym, arg_params)
diff --git a/src/c_api/c_api_symbolic.cc b/src/c_api/c_api_symbolic.cc
index 4666b6adf0c..e5e9b522890 100644
--- a/src/c_api/c_api_symbolic.cc
+++ b/src/c_api/c_api_symbolic.cc
@@ -577,7 +577,8 @@ int MXQuantizeSymbol(SymbolHandle sym_handle,
const mx_uint num_excluded_symbols,
const SymbolHandle *excluded_symbols,
const mx_uint num_offline,
- const char **offline_params) {
+ const char **offline_params,
+ const char *quantized_dtype) {
nnvm::Symbol *s = new nnvm::Symbol();
API_BEGIN();
nnvm::Symbol *sym = static_cast<nnvm::Symbol*>(sym_handle);
@@ -594,7 +595,9 @@ int MXQuantizeSymbol(SymbolHandle sym_handle,
for (size_t i = 0; i < num_offline; ++i) {
offline.emplace(offline_params[i]);
}
+ std::string quantized_type(quantized_dtype);
g.attrs["offline_params"] = std::make_shared<nnvm::any>(std::move(offline));
+ g.attrs["quantized_dtype"] =
std::make_shared<nnvm::any>(std::move(quantized_type));
g = ApplyPass(std::move(g), "QuantizeGraph");
s->outputs = g.outputs;
*ret_sym_handle = s;
diff --git a/src/operator/nn/convolution-inl.h
b/src/operator/nn/convolution-inl.h
index 5632d73c261..d40abaf1fd6 100644
--- a/src/operator/nn/convolution-inl.h
+++ b/src/operator/nn/convolution-inl.h
@@ -125,6 +125,8 @@ struct ConvolutionParam : public
dmlc::Parameter<ConvolutionParam> {
}
};
+void ConvolutionParamParser(nnvm::NodeAttrs* attrs);
+
typedef ParamOpSign<ConvolutionParam> ConvSignature;
} // namespace op
diff --git a/src/operator/nn/convolution.cc b/src/operator/nn/convolution.cc
index 0e8a929e1ba..ef70ccd6ec1 100644
--- a/src/operator/nn/convolution.cc
+++ b/src/operator/nn/convolution.cc
@@ -331,7 +331,7 @@ inline static bool BackwardConvStorageType(const
nnvm::NodeAttrs& attrs,
dispatch_mode, wanted_mode);
}
-static void ConvolutionParamParser(nnvm::NodeAttrs* attrs) {
+void ConvolutionParamParser(nnvm::NodeAttrs* attrs) {
using namespace mshadow;
ConvolutionParam param_;
try {
diff --git a/src/operator/nn/mkldnn/mkldnn_convolution-inl.h
b/src/operator/nn/mkldnn/mkldnn_convolution-inl.h
new file mode 100644
index 00000000000..23f2fe69463
--- /dev/null
+++ b/src/operator/nn/mkldnn/mkldnn_convolution-inl.h
@@ -0,0 +1,77 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file mkldnn_convolution-inl.h
+ * \brief
+*/
+
+#ifndef MXNET_OPERATOR_NN_MKLDNN_MKLDNN_CONVOLUTION_INL_H_
+#define MXNET_OPERATOR_NN_MKLDNN_MKLDNN_CONVOLUTION_INL_H_
+
+#if MXNET_USE_MKLDNN == 1
+
+#include <utility>
+#include "../convolution-inl.h"
+#include "./mkldnn_ops-inl.h"
+#include "./mkldnn_base-inl.h"
+
+namespace mxnet {
+namespace op {
+
+mkldnn::convolution_forward::primitive_desc GetConvFwdImpl(
+ const ConvolutionParam& param, const bool is_train, const NDArray &data,
+ const NDArray &weights, const NDArray *bias, const NDArray &output);
+
+class MKLDNNConvForward {
+ public:
+ mkldnn::convolution_forward::primitive_desc fwd_pd;
+
+ MKLDNNConvForward(const ConvolutionParam& param, const bool is_train,
+ const NDArray &data, const NDArray &weights,
+ const NDArray *bias, const NDArray &output): fwd_pd(
+ GetConvFwdImpl(param, is_train, data, weights, bias,
output)) {
+ }
+
+ void SetNewMem(const mkldnn::memory &data, const mkldnn::memory &weight,
+ const mkldnn::memory *bias, const mkldnn::memory &output);
+
+ const mkldnn::convolution_forward &GetFwd() const {
+ return *fwd_;
+ }
+
+ private:
+ std::shared_ptr<mkldnn::convolution_forward> fwd_;
+ std::shared_ptr<mkldnn::memory> data_;
+ std::shared_ptr<mkldnn::memory> weight_;
+ std::shared_ptr<mkldnn::memory> bias_;
+ std::shared_ptr<mkldnn::memory> out_;
+};
+
+typedef ParamOpSign<ConvolutionParam> MKLDNNConvSignature;
+
+MKLDNNConvForward &GetConvFwd(const nnvm::NodeAttrs& attrs,
+ const bool is_train, const NDArray &data, const NDArray &weights,
+ const NDArray *bias, const NDArray &output);
+
+} // namespace op
+} // namespace mxnet
+
+#endif // MXNET_USE_MKLDNN == 1
+#endif // MXNET_OPERATOR_NN_MKLDNN_MKLDNN_CONVOLUTION_INL_H_
diff --git a/src/operator/nn/mkldnn/mkldnn_convolution.cc
b/src/operator/nn/mkldnn/mkldnn_convolution.cc
index f851a6d2535..cf04ea8da3d 100644
--- a/src/operator/nn/mkldnn/mkldnn_convolution.cc
+++ b/src/operator/nn/mkldnn/mkldnn_convolution.cc
@@ -23,11 +23,14 @@
* \author Da Zheng
*/
+
+#if MXNET_USE_MKLDNN == 1
+
#include "../convolution-inl.h"
#include "./mkldnn_ops-inl.h"
#include "./mkldnn_base-inl.h"
+#include "./mkldnn_convolution-inl.h"
-#if MXNET_USE_MKLDNN == 1
namespace mxnet {
namespace op {
@@ -37,8 +40,8 @@ bool SupportMKLDNNConv(const ConvolutionParam& params, const
NDArray &input) {
return input.dtype() == mshadow::kFloat32 && input.shape().ndim() == 4;
}
-static mkldnn::convolution_forward::primitive_desc GetConvFwdImpl(
- const ConvolutionParam& param, bool is_train, const NDArray &data,
+mkldnn::convolution_forward::primitive_desc GetConvFwdImpl(
+ const ConvolutionParam& param, const bool is_train, const NDArray &data,
const NDArray &weights, const NDArray *bias, const NDArray &output) {
auto prop = is_train ? mkldnn::prop_kind::forward_training :
mkldnn::prop_kind::forward_scoring;
auto data_md = GetMemDesc(data);
@@ -162,73 +165,51 @@ static
mkldnn::convolution_backward_weights::primitive_desc GetConvBwdWeights(
}
}
-class MKLDNNConvForward {
- std::shared_ptr<mkldnn::convolution_forward> fwd;
- std::shared_ptr<mkldnn::memory> data;
- std::shared_ptr<mkldnn::memory> weight;
- std::shared_ptr<mkldnn::memory> bias;
- std::shared_ptr<mkldnn::memory> out;
+void MKLDNNConvForward::SetNewMem(const mkldnn::memory &data,
+ const mkldnn::memory &weight,
+ const mkldnn::memory *bias,
+ const mkldnn::memory &output) {
+ if (this->data_ == nullptr)
+ this->data_ = std::shared_ptr<mkldnn::memory>(new mkldnn::memory(
+ fwd_pd.src_primitive_desc(), data.get_data_handle()));
+ else
+ this->data_->set_data_handle(data.get_data_handle());
- public:
- mkldnn::convolution_forward::primitive_desc fwd_pd;
+ if (this->weight_ == nullptr)
+ this->weight_ = std::shared_ptr<mkldnn::memory>(new mkldnn::memory(
+ fwd_pd.weights_primitive_desc(), weight.get_data_handle()));
+ else
+ this->weight_->set_data_handle(weight.get_data_handle());
- MKLDNNConvForward(const ConvolutionParam& param, bool is_train,
- const NDArray &data, const NDArray &weights,
- const NDArray *bias, const NDArray &output): fwd_pd(
- GetConvFwdImpl(param, is_train, data, weights, bias,
output)) {
- }
-
- void SetNewMem(const mkldnn::memory &data, const mkldnn::memory &weight,
- const mkldnn::memory *bias, const mkldnn::memory &output) {
- if (this->data == nullptr)
- this->data = std::shared_ptr<mkldnn::memory>(new mkldnn::memory(
- fwd_pd.src_primitive_desc(), data.get_data_handle()));
- else
- this->data->set_data_handle(data.get_data_handle());
+ if (this->out_ == nullptr)
+ this->out_ = std::shared_ptr<mkldnn::memory>(new mkldnn::memory(
+ fwd_pd.dst_primitive_desc(), output.get_data_handle()));
+ else
+ this->out_->set_data_handle(output.get_data_handle());
- if (this->weight == nullptr)
- this->weight = std::shared_ptr<mkldnn::memory>(new mkldnn::memory(
- fwd_pd.weights_primitive_desc(), weight.get_data_handle()));
+ if (bias != nullptr) {
+ if (this->bias_ == nullptr)
+ this->bias_ = std::shared_ptr<mkldnn::memory>(new mkldnn::memory(
+ fwd_pd.bias_primitive_desc(), bias->get_data_handle()));
else
- this->weight->set_data_handle(weight.get_data_handle());
-
- if (this->out == nullptr)
- this->out = std::shared_ptr<mkldnn::memory>(new mkldnn::memory(
- fwd_pd.dst_primitive_desc(), output.get_data_handle()));
- else
- this->out->set_data_handle(output.get_data_handle());
-
- if (bias != nullptr) {
- if (this->bias == nullptr)
- this->bias = std::shared_ptr<mkldnn::memory>(new mkldnn::memory(
- fwd_pd.bias_primitive_desc(), bias->get_data_handle()));
- else
- this->bias->set_data_handle(bias->get_data_handle());
- if (this->fwd == nullptr)
- this->fwd = std::shared_ptr<mkldnn::convolution_forward>(
- new mkldnn::convolution_forward(fwd_pd,
mkldnn::primitive::at(*this->data),
-
mkldnn::primitive::at(*this->weight),
- mkldnn::primitive::at(*this->bias),
- *this->out));
- } else if (this->fwd == nullptr) {
- this->fwd = std::shared_ptr<mkldnn::convolution_forward>(
- new mkldnn::convolution_forward(fwd_pd,
mkldnn::primitive::at(*this->data),
- mkldnn::primitive::at(*this->weight),
- *this->out));
- }
+ this->bias_->set_data_handle(bias->get_data_handle());
+ if (this->fwd_ == nullptr)
+ this->fwd_ = std::shared_ptr<mkldnn::convolution_forward>(
+ new mkldnn::convolution_forward(fwd_pd,
mkldnn::primitive::at(*this->data_),
+
mkldnn::primitive::at(*this->weight_),
+ mkldnn::primitive::at(*this->bias_),
+ *this->out_));
+ } else if (this->fwd_ == nullptr) {
+ this->fwd_ = std::shared_ptr<mkldnn::convolution_forward>(
+ new mkldnn::convolution_forward(fwd_pd,
mkldnn::primitive::at(*this->data_),
+ mkldnn::primitive::at(*this->weight_),
+ *this->out_));
}
+}
- const mkldnn::convolution_forward &GetFwd() const {
- return *fwd;
- }
-};
-
-typedef ParamOpSign<ConvolutionParam> MKLDNNConvSignature;
-
-static inline MKLDNNConvForward &GetConvFwd(
- const nnvm::NodeAttrs& attrs, bool is_train,
- const NDArray &data, const NDArray &weights,
- const NDArray *bias, const NDArray &output) {
+MKLDNNConvForward &GetConvFwd(const nnvm::NodeAttrs& attrs, const bool
is_train,
+ const NDArray &data, const NDArray &weights,
+ const NDArray *bias, const NDArray &output) {
#if DMLC_CXX11_THREAD_LOCAL
static thread_local std::unordered_map<MKLDNNConvSignature,
MKLDNNConvForward, OpHash> fwds;
#else
diff --git a/src/operator/nn/mkldnn/mkldnn_pooling-inl.h
b/src/operator/nn/mkldnn/mkldnn_pooling-inl.h
index 4b6235ec446..691e1d371b5 100644
--- a/src/operator/nn/mkldnn/mkldnn_pooling-inl.h
+++ b/src/operator/nn/mkldnn/mkldnn_pooling-inl.h
@@ -119,6 +119,10 @@ void MKLDNNPoolingGradCompute(const OpContext &ctx, const
PoolingParam ¶m,
const NDArray &out_grad, const NDArray &in_data,
const NDArray *workspace, const OpReqType req,
const NDArray &in_grad);
+MKLDNNPoolingFwd &GetPoolingFwd(const PoolingParam ¶m,
+ const bool is_train,
+ const NDArray &data,
+ const NDArray &output);
} // namespace op
} // namespace mxnet
#endif // MXNET_USE_MKLDNN == 1
diff --git a/src/operator/nn/pooling-inl.h b/src/operator/nn/pooling-inl.h
index 5993bf5149d..7ca6d2583db 100644
--- a/src/operator/nn/pooling-inl.h
+++ b/src/operator/nn/pooling-inl.h
@@ -41,6 +41,8 @@
namespace mxnet {
namespace op {
+void PoolingParamParser(nnvm::NodeAttrs *attrs);
+
struct PoolingParam : public dmlc::Parameter<PoolingParam> {
TShape kernel;
TShape stride;
diff --git a/src/operator/nn/pooling.cc b/src/operator/nn/pooling.cc
index ca472b3ca1b..b50cfe41e49 100644
--- a/src/operator/nn/pooling.cc
+++ b/src/operator/nn/pooling.cc
@@ -35,7 +35,7 @@
namespace mxnet {
namespace op {
-static void PoolingParamParser(nnvm::NodeAttrs *attrs) {
+void PoolingParamParser(nnvm::NodeAttrs *attrs) {
using namespace mshadow;
PoolingParam param;
param.Init(attrs->dict);
diff --git a/src/operator/quantization/dequantize.cc
b/src/operator/quantization/dequantize.cc
index 92b808dd460..bbd79417676 100644
--- a/src/operator/quantization/dequantize.cc
+++ b/src/operator/quantization/dequantize.cc
@@ -23,11 +23,31 @@
* \brief
*/
#include "./dequantize-inl.h"
+#if MXNET_USE_MKLDNN == 1
+#include "./mkldnn/mkldnn_dequantize-inl.h"
+#endif
namespace mxnet {
namespace op {
DMLC_REGISTER_PARAMETER(DequantizeParam);
+bool DequantizeStorageType(const nnvm::NodeAttrs& attrs,
+ const int dev_mask,
+ DispatchMode* dispatch_mode,
+ std::vector<int> *in_attrs,
+ std::vector<int> *out_attrs) {
+ *dispatch_mode = DispatchMode::kFCompute;
+#if MXNET_USE_MKLDNN == 1
+ if (dev_mask == mshadow::cpu::kDevMask) {
+ *dispatch_mode = DispatchMode::kFComputeEx;
+ }
+#endif
+ (*out_attrs)[0] = kDefaultStorage;
+ (*out_attrs)[1] = kDefaultStorage;
+ (*out_attrs)[2] = kDefaultStorage;
+ return true;
+}
+
NNVM_REGISTER_OP(_contrib_dequantize)
.describe(R"code(Dequantize the input tensor into a float tensor.
min_range and max_range are scalar floats that specify the range for
@@ -50,6 +70,10 @@ by keep zero centered for the quantized value:
.set_num_outputs(1)
.set_attr<nnvm::FInferShape>("FInferShape", DequantizeShape)
.set_attr<nnvm::FInferType>("FInferType", DequantizeType)
+.set_attr<FInferStorageType>("FInferStorageType", DequantizeStorageType)
+#if MXNET_USE_MKLDNN == 1
+.set_attr<FComputeEx>("FComputeEx<cpu>", MKLDNNDequantizeCompute)
+#endif
.set_attr<FCompute>("FCompute<cpu>", DequantizeCompute<cpu>)
.add_argument("data", "NDArray-or-Symbol", "A ndarray/symbol of type `uint8`")
.add_argument("min_range", "NDArray-or-Symbol", "The minimum scalar value "
diff --git a/src/operator/quantization/mkldnn/mkldnn_dequantize-inl.h
b/src/operator/quantization/mkldnn/mkldnn_dequantize-inl.h
new file mode 100644
index 00000000000..89c3c199488
--- /dev/null
+++ b/src/operator/quantization/mkldnn/mkldnn_dequantize-inl.h
@@ -0,0 +1,105 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file mkldnn_dequantize-inl.h
+ * \author Wenting Jiang, Xinyu Chen
+ * \brief
+ */
+
+#ifndef MXNET_OPERATOR_QUANTIZATION_MKLDNN_MKLDNN_DEQUANTIZE_INL_H_
+#define MXNET_OPERATOR_QUANTIZATION_MKLDNN_MKLDNN_DEQUANTIZE_INL_H_
+#if MXNET_USE_MKLDNN == 1
+#include <string>
+#include <algorithm>
+#include <vector>
+#include "../../nn/mkldnn/mkldnn_base-inl.h"
+
+namespace mxnet {
+namespace op {
+
+template<typename SrcType, typename DstType>
+static void MKLDNNDequantizeComputeKer(const std::vector<NDArray> &inputs,
+ const std::vector<NDArray> &outputs,
+ const std::vector<OpReqType> &req) {
+ using namespace mshadow;
+ using namespace mxnet_op;
+ using red::limits::MaxValue;
+ using red::limits::MinValue;
+ float real_range = 0.0;
+ float quantized_range = 0.0;
+ if (inputs[0].dtype() == mshadow::kUint8) {
+ quantized_range = MaxAbs(MaxValue<SrcType>(), MinValue<SrcType>());
+ real_range = MaxAbs(*inputs[1].data().dptr<DstType>(),
*inputs[2].data().dptr<DstType>());
+ } else if (inputs[0].dtype() == mshadow::kInt8) {
+ quantized_range = MinAbs(MaxValue<SrcType>(), MinValue<SrcType>());
+ real_range = MaxAbs(*inputs[1].data().dptr<DstType>(),
*inputs[2].data().dptr<DstType>());
+ } else {
+ LOG(FATAL) << "mkldnn dequantize op only supports int8 and uint8 as output
type";
+ }
+ float scale = real_range / quantized_range;
+ primitive_attr attr;
+ const int mask = 0;
+ std::vector<float> scales = {scale};
+ attr.set_output_scales(mask, scales);
+ attr.set_int_output_round_mode(round_nearest);
+ mkldnn::engine cpu_engine = mxnet::CpuEngine::Get()->get_engine();
+
+ NDArray in_buffer = inputs[0];
+ if (inputs[0].IsView() && inputs[0].IsMKLDNNData())
+ in_buffer = inputs[0].Reorder2Default();
+
+ auto i_mem = in_buffer.GetMKLDNNData();
+ auto i_mpd = i_mem->get_primitive_desc();
+ auto i_desc = i_mpd.desc();
+ size_t i_ndim = in_buffer.shape().ndim();
+ mkldnn::memory::dims i_dims = mkldnn::memory::dims(i_ndim);
+ for (size_t i = 0; i < i_ndim; i++) {
+ i_dims[i] = static_cast<int>(in_buffer.shape()[i]);
+ }
+ mkldnn::memory::format i_fmt =
static_cast<mkldnn::memory::format>(i_desc.data.format);
+ auto o_desc = mkldnn::memory::desc(i_dims,
+
(mkldnn::memory::data_type)data_type_enum<DstType>::type,
+ i_fmt);
+ auto o_mpd = memory::primitive_desc(o_desc, cpu_engine);
+ auto reorder_pd = reorder::primitive_desc(i_mpd, o_mpd, attr);
+ auto o_mem = CreateMKLDNNMem(outputs[0], o_mpd, req[0]);
+ MKLDNNStream::Get()->RegisterPrim(mkldnn::reorder(reorder_pd, *i_mem,
*o_mem.second));
+ CommitOutput(outputs[0], o_mem);
+ MKLDNNStream::Get()->Submit();
+}
+
+static void MKLDNNDequantizeCompute(const nnvm::NodeAttrs& attrs, const
OpContext &ctx,
+ const std::vector<NDArray> &inputs,
+ const std::vector<OpReqType> &req,
+ const std::vector<NDArray> &outputs) {
+ if (inputs[0].dtype() == mshadow::kUint8) {
+ MKLDNNDequantizeComputeKer<uint8_t, float>(inputs, outputs, req);
+ } else if (inputs[0].dtype() == mshadow::kInt8) {
+ MKLDNNDequantizeComputeKer<int8_t, float>(inputs, outputs, req);
+ } else {
+ LOG(FATAL) << "mkldnn dequantize op only supports int8 and uint8 as input
type";
+ }
+}
+
+} // namespace op
+} // namespace mxnet
+
+#endif // MXNET_USE_MKLDNN == 1
+#endif // MXNET_OPERATOR_QUANTIZATION_MKLDNN_MKLDNN_DEQUANTIZE_INL_H_
diff --git a/src/operator/quantization/mkldnn/mkldnn_quantize-inl.h
b/src/operator/quantization/mkldnn/mkldnn_quantize-inl.h
new file mode 100644
index 00000000000..f7709319d6a
--- /dev/null
+++ b/src/operator/quantization/mkldnn/mkldnn_quantize-inl.h
@@ -0,0 +1,112 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file mkldnn_quantize-inl.h
+ * \brief
+ * \author Wenting Jiang, Xinyu Chen
+ */
+
+#ifndef MXNET_OPERATOR_QUANTIZATION_MKLDNN_MKLDNN_QUANTIZE_INL_H_
+#define MXNET_OPERATOR_QUANTIZATION_MKLDNN_MKLDNN_QUANTIZE_INL_H_
+#if MXNET_USE_MKLDNN == 1
+#include <string>
+#include <algorithm>
+#include <vector>
+#include "../quantize-inl.h"
+#include "../../nn/mkldnn/mkldnn_base-inl.h"
+
+namespace mxnet {
+namespace op {
+
+template<typename SrcType, typename DstType>
+static void MKLDNNQuantizeComputeKer(const std::vector<NDArray>& inputs,
+ const std::vector<NDArray>& outputs,
+ const QuantizeParam& param,
+ const std::vector<OpReqType> &req) {
+ using namespace mshadow;
+ using namespace mxnet_op;
+ using red::limits::MaxValue;
+ using red::limits::MinValue;
+ float real_range = 0.0;
+ float quantized_range = 0.0;
+ if (param.out_type == mshadow::kUint8) {
+ real_range = MaxAbs(*inputs[1].data().dptr<float>(),
*inputs[2].data().dptr<float>());
+ quantized_range = MaxAbs(MaxValue<DstType>(), MinValue<DstType>());
+ *outputs[1].data().dptr<float>() = *inputs[1].data().dptr<float>();
+ *outputs[2].data().dptr<float>() = *inputs[2].data().dptr<float>();
+ } else if (param.out_type == mshadow::kInt8) {
+ real_range = MaxAbs(*inputs[1].data().dptr<float>(),
*inputs[2].data().dptr<float>());
+ quantized_range = MinAbs(MaxValue<DstType>(), MinValue<DstType>());
+ *outputs[1].data().dptr<float>() = -real_range;
+ *outputs[2].data().dptr<float>() = real_range;
+ } else {
+ LOG(FATAL) << "mkldnn quantize op only supports int8 and uint8 as output
type";
+ }
+ float scale = quantized_range / real_range;
+ primitive_attr attr;
+ const int mask = 0;
+ std::vector<float> scales = {scale};
+ attr.set_output_scales(mask, scales);
+ attr.set_int_output_round_mode(round_nearest);
+ mkldnn::engine cpu_engine = mxnet::CpuEngine::Get()->get_engine();
+
+ NDArray in_buffer = inputs[0];
+ if (inputs[0].IsView() && inputs[0].IsMKLDNNData())
+ in_buffer = inputs[0].Reorder2Default();
+
+ auto i_mem = in_buffer.GetMKLDNNData();
+ auto i_mpd = i_mem->get_primitive_desc();
+ auto i_desc = i_mpd.desc();
+ mkldnn::memory::format i_fmt =
static_cast<mkldnn::memory::format>(i_desc.data.format);
+ size_t i_ndim = in_buffer.shape().ndim();
+ mkldnn::memory::dims i_dims = mkldnn::memory::dims(i_ndim);
+ for (size_t i = 0; i < i_ndim; i++) {
+ i_dims[i] = static_cast<int>(in_buffer.shape()[i]);
+ }
+ auto o_desc = mkldnn::memory::desc(i_dims,
+
(mkldnn::memory::data_type)data_type_enum<DstType>::type,
+ i_fmt);
+ auto o_mpd = memory::primitive_desc(o_desc, cpu_engine);
+ auto reorder_pd = reorder::primitive_desc(i_mpd, o_mpd, attr);
+ auto o_mem = CreateMKLDNNMem(outputs[0], o_mpd, req[0]);
+ MKLDNNStream::Get()->RegisterPrim(mkldnn::reorder(reorder_pd, *i_mem,
*o_mem.second));
+ CommitOutput(outputs[0], o_mem);
+ MKLDNNStream::Get()->Submit();
+}
+
+static void MKLDNNQuantizeCompute(const nnvm::NodeAttrs& attrs, const
OpContext &ctx,
+ const std::vector<NDArray> &inputs,
+ const std::vector<OpReqType> &req,
+ const std::vector<NDArray> &outputs) {
+ const QuantizeParam& param = nnvm::get<QuantizeParam>(attrs.parsed);
+ if (param.out_type == mshadow::kUint8) {
+ MKLDNNQuantizeComputeKer<float, uint8_t>(inputs, outputs, param, req);
+ } else if (param.out_type == mshadow::kInt8) {
+ MKLDNNQuantizeComputeKer<float, int8_t>(inputs, outputs, param, req);
+ } else {
+ LOG(FATAL) << "mkldnn quantize op only supports int8 and uint8 as output
type";
+ }
+}
+
+} // namespace op
+} // namespace mxnet
+
+#endif // MXNET_USE_MKLDNN == 1
+#endif // MXNET_OPERATOR_QUANTIZATION_MKLDNN_MKLDNN_QUANTIZE_INL_H_
diff --git a/src/operator/quantization/mkldnn/mkldnn_quantized_conv.cc
b/src/operator/quantization/mkldnn/mkldnn_quantized_conv.cc
new file mode 100644
index 00000000000..fa6a32a4739
--- /dev/null
+++ b/src/operator/quantization/mkldnn/mkldnn_quantized_conv.cc
@@ -0,0 +1,89 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file mkldnn_quantized_conv.cc
+ * \brief
+ * \author Wenting Jiang, Xinyu Chen
+*/
+
+#if MXNET_USE_MKLDNN == 1
+#include "../../nn/mkldnn/mkldnn_base-inl.h"
+#include "../../nn/mkldnn/mkldnn_convolution-inl.h"
+#include "../../nn/convolution-inl.h"
+#include "../quantization_utils.h"
+#include "../../tensor/matrix_op-inl.h"
+#include "../../elemwise_op_common.h"
+namespace mxnet {
+namespace op {
+
+static void MKLDNNQuantizedConvForward(const nnvm::NodeAttrs& attrs,
+ const OpContext &ctx,
+ const std::vector<NDArray> &in_data,
+ const std::vector<OpReqType> &req,
+ const std::vector<NDArray> &out_data) {
+ CHECK_EQ(in_data[0].dtype(), mshadow::kUint8)
+ << "mkldnn_quantized_conv op only supports uint8 as input type";
+ TmpMemMgr::Get()->Init(ctx.requested[conv::kTempSpace]);
+ const ConvolutionParam& param = nnvm::get<ConvolutionParam>(attrs.parsed);
+ NDArray weight = in_data[conv::kWeight];
+ MKLDNNConvForward &fwd = GetConvFwd(attrs, ctx.is_train,
+ in_data[conv::kData], weight,
+ param.no_bias ? nullptr : &in_data[conv::kBias], out_data[conv::kOut]);
+
+ auto data_mem =
in_data[conv::kData].GetMKLDNNDataReorder(fwd.fwd_pd.src_primitive_desc());
+ const mkldnn::memory *weight_mem;
+ // For inference, we want to reorder the weight array so we don't need to
+ // reorder data every time.
+ if (weight.IsDefaultData()) {
+ weight_mem = GetWeights(weight, fwd.fwd_pd.weights_primitive_desc(),
param.num_group);
+ // We also need to modify the layout on the original weight array. The
+ // data conversion happens after the weight array is used.
+ weight.MKLDNNDataReorderAsync(fwd.fwd_pd.weights_primitive_desc());
+ } else {
+ weight_mem = weight.GetMKLDNNData();
+ CHECK(weight_mem->get_primitive_desc() ==
fwd.fwd_pd.weights_primitive_desc());
+ }
+ auto out_mem = CreateMKLDNNMem(out_data[conv::kOut],
fwd.fwd_pd.dst_primitive_desc(),
+ req[conv::kOut]);
+ const mkldnn::memory *bias_mem = nullptr;
+ if (!param.no_bias)
+ bias_mem =
in_data[conv::kBias].GetMKLDNNDataReorder(fwd.fwd_pd.bias_primitive_desc());
+ fwd.SetNewMem(*data_mem, *weight_mem, bias_mem, *out_mem.second);
+ MKLDNNStream::Get()->RegisterPrim(fwd.GetFwd());
+
+ CommitOutput(out_data[conv::kOut], out_mem);
+ MKLDNNStream::Get()->Submit();
+ Stream<cpu> *s = ctx.get_stream<cpu>();
+ const size_t num_inputs = param.no_bias ? 2 : 3;
+ mxnet_op::Kernel<QuantizationRangeForMultiplicationStruct, cpu>::Launch(s, 1,
+ out_data[1].data().dptr<float>(), out_data[2].data().dptr<float>(),
+ in_data[num_inputs].data().dptr<float>(),
+ in_data[num_inputs+1].data().dptr<float>(),
+ in_data[num_inputs+2].data().dptr<float>(),
+ in_data[num_inputs+3].data().dptr<float>());
+}
+
+NNVM_REGISTER_OP(_contrib_quantized_conv)
+.set_attr<FComputeEx>("FComputeEx<cpu>", MKLDNNQuantizedConvForward);
+
+} // namespace op
+} // namespace mxnet
+
+#endif // MXNET_USE_MKLDNN == 1
diff --git a/src/operator/quantization/mkldnn/mkldnn_quantized_pooling.cc
b/src/operator/quantization/mkldnn/mkldnn_quantized_pooling.cc
new file mode 100644
index 00000000000..83177ad9b34
--- /dev/null
+++ b/src/operator/quantization/mkldnn/mkldnn_quantized_pooling.cc
@@ -0,0 +1,54 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file mkldnn_quantized_pooling.cc
+ * \brief
+ * \author Tao Lv, Xinyu Chen
+*/
+
+#if MXNET_USE_MKLDNN == 1
+
+#include "../../nn/mkldnn/mkldnn_pooling-inl.h"
+
+namespace mxnet {
+namespace op {
+
+static void MKLDNNQuantizedPoolingForward(const nnvm::NodeAttrs& attrs, const
OpContext &ctx,
+ const std::vector<NDArray> &in_data,
+ const std::vector<OpReqType> &req,
+ const std::vector<NDArray>
&out_data) {
+ CHECK(in_data[0].dtype() == mshadow::kUint8
+ || in_data[0].dtype() == mshadow::kInt8)
+ << "mkldnn_quantized_pooling op only supports uint8 and int8 as input
type";
+ const PoolingParam& param = nnvm::get<PoolingParam>(attrs.parsed);
+ auto fwd = GetPoolingFwd(param, ctx.is_train, in_data[0], out_data[0]);
+ fwd.SetDataHandle(in_data[0], out_data[0]);
+ fwd.Execute();
+ out_data[1].data().dptr<float>()[0] = in_data[1].data().dptr<float>()[0];
+ out_data[2].data().dptr<float>()[0] = in_data[2].data().dptr<float>()[0];
+}
+
+NNVM_REGISTER_OP(_contrib_quantized_pooling)
+.set_attr<FComputeEx>("FComputeEx<cpu>", MKLDNNQuantizedPoolingForward);
+
+} // namespace op
+} // namespace mxnet
+
+#endif // MXNET_USE_MKLDNN == 1
diff --git a/src/operator/quantization/mkldnn/mkldnn_requantize-inl.h
b/src/operator/quantization/mkldnn/mkldnn_requantize-inl.h
new file mode 100644
index 00000000000..409c53dd3b9
--- /dev/null
+++ b/src/operator/quantization/mkldnn/mkldnn_requantize-inl.h
@@ -0,0 +1,158 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/* \file mkldnn_requantize-inl.h
+ * \brief
+ * \author Jin Huang, Xinyu Chen
+ */
+
+#ifndef MXNET_OPERATOR_QUANTIZATION_MKLDNN_MKLDNN_REQUANTIZE_INL_H_
+#define MXNET_OPERATOR_QUANTIZATION_MKLDNN_MKLDNN_REQUANTIZE_INL_H_
+#if MXNET_USE_MKLDNN == 1
+#include <string>
+#include <algorithm>
+#include <vector>
+#include "../requantize-inl.h"
+#include "../../nn/mkldnn/mkldnn_base-inl.h"
+
+namespace mxnet {
+namespace op {
+
+static void MKLDNNRequantizeForwardKer(const nnvm::NodeAttrs& attrs,
+ const OpContext& ctx,
+ const std::vector<NDArray>& inputs,
+ const std::vector<OpReqType>& req,
+ const std::vector<NDArray>& outputs,
+ const float real_range) {
+ using namespace mshadow;
+ using namespace mxnet_op;
+ using red::limits::MaxValue;
+ using red::limits::MinValue;
+ typedef int32_t SrcDType;
+ typedef int8_t DstDType;
+ // check shapes
+ size_t i_dim = inputs[0].shape().ndim();
+ size_t o_dim = outputs[0].shape().ndim();
+ CHECK_EQ(i_dim, o_dim);
+ float first_quantized_range = MinAbs(MinValue<SrcDType>(),
+ MaxValue<SrcDType>());
+ float first_real_range = MaxAbs(*inputs[1].data().dptr<float>(),
+ *inputs[2].data().dptr<float>());
+ float first_scale = first_real_range / first_quantized_range;
+ float second_real_range = real_range;
+ float second_quantized_range = MinAbs(MaxValue<DstDType>(),
+ MinValue<DstDType>());
+ float second_scale = second_quantized_range / second_real_range;
+ float scale = first_scale * second_scale;
+ *outputs[1].data().dptr<float>() = -second_real_range;
+ *outputs[2].data().dptr<float>() = second_real_range;
+ primitive_attr attr;
+ const int mask = 0;
+ std::vector<float> scales = {scale};
+ attr.set_output_scales(mask, scales);
+ attr.set_int_output_round_mode(round_nearest);
+ mkldnn::engine cpu_engine = mxnet::CpuEngine::Get()->get_engine();
+
+ NDArray in_buffer = inputs[0];
+ if (inputs[0].IsView() && inputs[0].IsMKLDNNData())
+ in_buffer = inputs[0].Reorder2Default();
+
+ auto i_mem = in_buffer.GetMKLDNNData();
+ auto i_mpd = i_mem->get_primitive_desc();
+ auto i_desc = i_mpd.desc();
+ mkldnn::memory::format i_fmt =
static_cast<mkldnn::memory::format>(i_desc.data.format);
+ mkldnn::memory::dims i_dims = mkldnn::memory::dims(i_dim);
+ for (size_t i = 0; i < i_dim; i++) {
+ i_dims[i] = static_cast<int>(in_buffer.shape()[i]);
+ }
+ auto o_desc = mkldnn::memory::desc(i_dims,
+
(mkldnn::memory::data_type)data_type_enum<DstDType>::type,
+ i_fmt);
+ auto o_mpd = memory::primitive_desc(o_desc, cpu_engine);
+ auto reorder_pd = reorder::primitive_desc(i_mpd, o_mpd, attr);
+ auto o_mem = CreateMKLDNNMem(outputs[0], o_mpd, req[0]);
+ MKLDNNStream::Get()->RegisterPrim(mkldnn::reorder(reorder_pd, *i_mem,
*o_mem.second));
+ CommitOutput(outputs[0], o_mem);
+ MKLDNNStream::Get()->Submit();
+}
+
+static void MKLDNNRequantizeForward(const nnvm::NodeAttrs& attrs,
+ const OpContext& ctx,
+ const std::vector<NDArray>& inputs,
+ const std::vector<OpReqType>& req,
+ const std::vector<NDArray>& outputs) {
+ using namespace mshadow;
+ using namespace mxnet_op;
+ typedef int32_t SrcDType;
+ typedef int8_t DstDType;
+ Stream<cpu> *s = ctx.get_stream<cpu>();
+ const RequantizeParam& param = nnvm::get<RequantizeParam>(attrs.parsed);
+ float real_range;
+ // Model is calibrated
+ if (param.min_calib_range.has_value() && param.max_calib_range.has_value()) {
+ real_range =
+ MaxAbs(param.min_calib_range.value(), param.max_calib_range.value());
+ MKLDNNRequantizeForwardKer(attrs, ctx, inputs, req, outputs, real_range);
+ // Model is not calibrated
+ } else {
+ TShape src_shape, dst_shape;
+ const size_t actual_float_size = sizeof(float);
+ const size_t actual_quantized_size = sizeof(SrcDType);
+ const size_t temp_reduce_size = ConfigReduce<cpu, SrcDType>(s,
+ inputs[0].shape(), TShape({1}), &src_shape,
&dst_shape);
+ Tensor<cpu, 1, char> temp_space =
+ ctx.requested[0].get_space_typed<cpu, 1, char>(
+ Shape1(2*actual_float_size+2*actual_quantized_size+temp_reduce_size), s);
+ Tensor<cpu, 1, float> actual_min_float(
+ reinterpret_cast<float*>(temp_space.dptr_), Shape1(1), s);
+ Tensor<cpu, 1, float> actual_max_float(
+ reinterpret_cast<float*>(temp_space.dptr_) + 1, Shape1(1), s);
+ const int dev_id = ctx.run_ctx.ctx.dev_id;
+ TBlob actual_min_quantized(reinterpret_cast<SrcDType*>(
+ temp_space.dptr_ + 8), Shape1(1), cpu::kDevMask,
dev_id);
+ TBlob actual_max_quantized(reinterpret_cast<SrcDType*>(
+ temp_space.dptr_ + 8) + 1, Shape1(1), cpu::kDevMask,
dev_id);
+ Tensor<cpu, 1, char> workspace(
+ temp_space.dptr_+2*actual_float_size+2*actual_quantized_size,
+ Shape1(temp_reduce_size), s);
+ broadcast::Reduce<red::minimum, 2, SrcDType, mshadow::op::identity>(
+ s, actual_min_quantized.reshape(dst_shape), kWriteTo,
+ workspace, inputs[0].Reorder2Default().data().reshape(src_shape));
+ Kernel<QuantizedToFloatStruct, cpu>::Launch(s, 1,
+ actual_min_float.dptr_, actual_min_quantized.dptr<SrcDType>(),
+ inputs[1].Reorder2Default().data().dptr<float>(),
+ inputs[2].Reorder2Default().data().dptr<float>());
+ broadcast::Reduce<red::maximum, 2, SrcDType, mshadow::op::identity>(
+ s, actual_max_quantized.reshape(dst_shape), kWriteTo,
+ workspace, inputs[0].Reorder2Default().data().reshape(src_shape));
+ Kernel<QuantizedToFloatStruct, cpu>::Launch(s, 1,
+ actual_max_float.dptr_, actual_max_quantized.dptr<SrcDType>(),
+ inputs[1].Reorder2Default().data().dptr<float>(),
+ inputs[2].Reorder2Default().data().dptr<float>());
+
+ real_range = MaxAbs(*actual_min_float.dptr_, *actual_max_float.dptr_);
+ MKLDNNRequantizeForwardKer(attrs, ctx, inputs, req, outputs, real_range);
+ }
+}
+
+} // namespace op
+} // namespace mxnet
+
+#endif // MXNET_USE_MKLDNN == 1
+#endif // MXNET_OPERATOR_QUANTIZATION_MKLDNN_MKLDNN_REQUANTIZE_INL_H_
diff --git a/src/operator/quantization/quantize.cc
b/src/operator/quantization/quantize.cc
index 32eb952fa5d..25fb19dddd1 100644
--- a/src/operator/quantization/quantize.cc
+++ b/src/operator/quantization/quantize.cc
@@ -23,11 +23,31 @@
* \brief
*/
#include "./quantize-inl.h"
+#if MXNET_USE_MKLDNN == 1
+#include "./mkldnn/mkldnn_quantize-inl.h"
+#endif
namespace mxnet {
namespace op {
DMLC_REGISTER_PARAMETER(QuantizeParam);
+bool QuantizeStorageType(const nnvm::NodeAttrs& attrs,
+ const int dev_mask,
+ DispatchMode* dispatch_mode,
+ std::vector<int> *in_attrs,
+ std::vector<int> *out_attrs) {
+ *dispatch_mode = DispatchMode::kFCompute;
+#if MXNET_USE_MKLDNN == 1
+ if (dev_mask == mshadow::cpu::kDevMask) {
+ *dispatch_mode = DispatchMode::kFComputeEx;
+ }
+#endif
+ (*out_attrs)[0] = kDefaultStorage;
+ (*out_attrs)[1] = kDefaultStorage;
+ (*out_attrs)[2] = kDefaultStorage;
+ return true;
+}
+
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`.
@@ -61,6 +81,10 @@ where
})
.set_attr<nnvm::FInferShape>("FInferShape", QuantizeShape)
.set_attr<nnvm::FInferType>("FInferType", QuantizeType)
+.set_attr<FInferStorageType>("FInferStorageType", QuantizeStorageType)
+#if MXNET_USE_MKLDNN == 1
+.set_attr<FComputeEx>("FComputeEx<cpu>", MKLDNNQuantizeCompute)
+#endif
.set_attr<FCompute>("FCompute<cpu>", QuantizeCompute<cpu>)
.add_argument("data", "NDArray-or-Symbol", "A ndarray/symbol of type
`float32`")
.add_argument("min_range", "NDArray-or-Symbol", "The minimum scalar value "
diff --git a/src/operator/quantization/quantize_graph_pass.cc
b/src/operator/quantization/quantize_graph_pass.cc
index 5ec745ccdf3..5376a0ee9f1 100644
--- a/src/operator/quantization/quantize_graph_pass.cc
+++ b/src/operator/quantization/quantize_graph_pass.cc
@@ -99,6 +99,7 @@ Graph QuantizeGraph(Graph &&src) {
static auto& need_requantize_map =
Op::GetAttr<mxnet::FNeedRequantize>("FNeedRequantize");
auto offline_params =
src.GetAttr<std::unordered_set<std::string>>("offline_params");
auto excluded_nodes =
src.GetAttr<std::unordered_set<NodePtr>>("excluded_nodes");
+ auto quantized_dtype = src.GetAttr<std::string>("quantized_dtype");
// mirror_map stores the mapping from the currently visited graph to the
newly created quantized
// graph. Key is the currently visited graph's node pointer, and value is a
copied node of the key
@@ -129,7 +130,7 @@ Graph QuantizeGraph(Graph &&src) {
mirror_node->op()->name != "_contrib_quantize")) {
NodePtr quantize_node = InsertNode("_contrib_quantize",
e.node->attrs.name + "_quantize", new_node, mirror_entry);
- quantize_node->attrs.dict["out_type"] = "int8";
+ quantize_node->attrs.dict["out_type"] = quantized_dtype;
quantize_node->op()->attr_parser(&(quantize_node->attrs));
NodePtr min_node = InsertNode("min",
@@ -159,7 +160,11 @@ Graph QuantizeGraph(Graph &&src) {
uint32_t min_index = 1;
uint32_t max_index = 2;
if (quantized_op_map.count(e.node->op())) {
- size_t num_outputs = e.node->num_outputs();
+ // here we calculate the output number (exclude min/max, in order to
+ // calculate min/max index from mirror node) based on assumption that
+ // there is only 1min and 1max output from mirror node (which is
+ // currently true)
+ size_t num_outputs = mirror_node->num_outputs() - 2;
min_index = num_outputs + 2 * e.index;
max_index = num_outputs + 2 * e.index + 1;
} else {
@@ -198,12 +203,15 @@ Graph QuantizeGraph(Graph &&src) {
NodePtr mirror_node = mirror_map.at(e.node.get());
NodeEntry mirror_entry = NodeEntry{
mirror_node, e.index, e.version};
- size_t num_outputs = e.node->num_outputs();
- uint32_t min_index = num_outputs + 2 * e.index;
- uint32_t max_index = num_outputs + 2 * e.index + 1;
-
// if input node is quantized operator, add dequantize node
if (NeedQuantize(e.node, excluded_nodes)) {
+ // here we calculate the output number (exclude min/max, in order to
+ // calculate min/max index from mirror node) based on assumption that
+ // there is only 1min and 1max output from mirror node (which is
+ // currently true)
+ size_t num_outputs = mirror_node->num_outputs() - 2;
+ uint32_t min_index = num_outputs + 2 * e.index;
+ uint32_t max_index = num_outputs + 2 * e.index + 1;
NodePtr dequantize_node = CreateNode("_contrib_dequantize",
e.node->attrs.name + "_dequantize");
dequantize_node->inputs.emplace_back(mirror_entry);
diff --git a/src/operator/quantization/quantized_conv.cc
b/src/operator/quantization/quantized_conv.cc
index d7dc9fe4dbd..ed62228b924 100644
--- a/src/operator/quantization/quantized_conv.cc
+++ b/src/operator/quantization/quantized_conv.cc
@@ -24,6 +24,9 @@
* \author Ziheng Jiang, Jun Wu
*/
#include "../nn/convolution-inl.h"
+#if MXNET_USE_MKLDNN == 1
+#include "../nn/mkldnn/mkldnn_ops-inl.h"
+#endif
namespace mxnet {
namespace op {
@@ -86,12 +89,13 @@ bool QuantizedConvType(const nnvm::NodeAttrs& attrs,
const ConvolutionParam& param = nnvm::get<ConvolutionParam>(attrs.parsed);
CHECK_EQ(in_type->size(), param.no_bias? 6U : 9U);
CHECK_EQ(out_type->size(), 3U);
+#ifndef MXNET_USE_MKLDNN
TYPE_ASSIGN_CHECK(*in_type, 0, mshadow::kInt8);
+#endif
TYPE_ASSIGN_CHECK(*in_type, 1, mshadow::kInt8);
if (!param.no_bias) {
TYPE_ASSIGN_CHECK(*in_type, 2, mshadow::kInt8);
}
-
const size_t start = param.no_bias? 2 : 3;
const size_t end = param.no_bias? 6 : 9;
for (size_t i = start; i < end; ++i) {
@@ -104,6 +108,24 @@ bool QuantizedConvType(const nnvm::NodeAttrs& attrs,
return true;
}
+bool QuantizedConvStorageType(const nnvm::NodeAttrs& attrs,
+ const int dev_mask,
+ DispatchMode* dispatch_mode,
+ std::vector<int> *in_attrs,
+ std::vector<int> *out_attrs) {
+ *dispatch_mode = DispatchMode::kFCompute;
+#if MXNET_USE_MKLDNN == 1
+ if (dev_mask == mshadow::cpu::kDevMask) {
+ *dispatch_mode = DispatchMode::kFComputeEx;
+ }
+#endif
+
+ (*out_attrs)[0] = kDefaultStorage;
+ (*out_attrs)[1] = kDefaultStorage;
+ (*out_attrs)[2] = kDefaultStorage;
+ return true;
+}
+
NNVM_REGISTER_OP(_contrib_quantized_conv)
.describe(R"code(Convolution operator for input, weight and bias data type of
int8,
and accumulates in type int32 for the output. For each argument, two more
arguments of type
@@ -119,7 +141,7 @@ and max thresholds representing the threholds for
quantizing the float32 output
return param.no_bias? 6 : 9;
})
.set_num_outputs(3)
-.set_attr_parser(ParamParser<ConvolutionParam>)
+.set_attr_parser(ConvolutionParamParser)
.set_attr<nnvm::FListInputNames>("FListInputNames",
[](const NodeAttrs& attrs) {
const ConvolutionParam& param = nnvm::get<ConvolutionParam>(attrs.parsed);
@@ -137,6 +159,7 @@ and max thresholds representing the threholds for
quantizing the float32 output
})
.set_attr<nnvm::FInferShape>("FInferShape", QuantizedConvShape)
.set_attr<nnvm::FInferType>("FInferType", QuantizedConvType)
+.set_attr<FInferStorageType>("FInferStorageType", QuantizedConvStorageType)
.set_attr<FResourceRequest>("FResourceRequest",
[](const NodeAttrs& attrs) {
return std::vector<ResourceRequest>(1, ResourceRequest::kTempSpace);
diff --git a/src/operator/quantization/quantized_flatten-inl.h
b/src/operator/quantization/quantized_flatten-inl.h
index 95f36615402..b7209fd28f5 100644
--- a/src/operator/quantization/quantized_flatten-inl.h
+++ b/src/operator/quantization/quantized_flatten-inl.h
@@ -62,11 +62,21 @@ void QuantizedFlattenCompute(const nnvm::NodeAttrs& attrs,
using namespace mxnet_op;
Stream<xpu> *s = ctx.get_stream<xpu>();
- typedef int8_t DstDType;
- typedef int8_t SrcDType;
- Kernel<quantized_flatten, 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>());
+ if (inputs[0].type_flag_ == mshadow::kUint8) {
+ typedef uint8_t SrcDType;
+ typedef uint8_t DstDType;
+ Kernel<quantized_flatten, 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>());
+ } else if (inputs[0].type_flag_ == mshadow::kInt8) {
+ typedef int8_t SrcDType;
+ typedef int8_t DstDType;
+ Kernel<quantized_flatten, 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>());
+ } else {
+ LOG(FATAL) << "quantized_flatten op only supports int8 and uint8 as input
and output type";
+ }
}
inline bool QuantizedFlattenShape(const nnvm::NodeAttrs& attrs,
@@ -96,10 +106,9 @@ inline bool QuantizedFlattenType(const nnvm::NodeAttrs&
attrs,
std::vector<int> *out_attrs) {
CHECK_EQ(in_attrs->size(), 3U);
CHECK_EQ(out_attrs->size(), 3U);
- TYPE_ASSIGN_CHECK(*in_attrs, 0, mshadow::kInt8);
TYPE_ASSIGN_CHECK(*in_attrs, 1, mshadow::kFloat32);
TYPE_ASSIGN_CHECK(*in_attrs, 2, mshadow::kFloat32);
- TYPE_ASSIGN_CHECK(*out_attrs, 0, mshadow::kInt8);
+ TYPE_ASSIGN_CHECK(*out_attrs, 0, (*in_attrs)[0]);
TYPE_ASSIGN_CHECK(*out_attrs, 1, mshadow::kFloat32);
TYPE_ASSIGN_CHECK(*out_attrs, 2, mshadow::kFloat32);
return (*in_attrs)[0] != -1;
diff --git a/src/operator/quantization/quantized_pooling.cc
b/src/operator/quantization/quantized_pooling.cc
index a3105eb654d..779e244c862 100644
--- a/src/operator/quantization/quantized_pooling.cc
+++ b/src/operator/quantization/quantized_pooling.cc
@@ -23,6 +23,9 @@
*/
#include <mxnet/op_attr_types.h>
#include "../nn/pooling-inl.h"
+#if MXNET_USE_MKLDNN == 1
+#include "../nn/mkldnn/mkldnn_pooling-inl.h"
+#endif
namespace mxnet {
namespace op {
@@ -79,8 +82,12 @@ bool QuantizedPoolingType(const nnvm::NodeAttrs& attrs,
CHECK_EQ(in_type->size(), 3U);
CHECK_EQ(out_type->size(), 3U);
if (param.pool_type == pool_enum::kMaxPooling || param.pool_type ==
pool_enum::kAvgPooling) {
+#if MXNET_USE_MKLDNN == 1
+ TYPE_ASSIGN_CHECK(*out_type, 0, (*in_type)[0]);
+#else
TYPE_ASSIGN_CHECK(*in_type, 0, mshadow::kInt8);
TYPE_ASSIGN_CHECK(*out_type, 0, mshadow::kInt8);
+#endif
} else {
LOG(FATAL) << "QuantizedPoolingOp only supports pool_type=max/avg for now";
}
@@ -91,6 +98,27 @@ bool QuantizedPoolingType(const nnvm::NodeAttrs& attrs,
return true;
}
+inline static bool QuantizedPoolingStorageType(const nnvm::NodeAttrs &attrs,
+ const int dev_mask,
+ DispatchMode *dispatch_mode,
+ std::vector<int> *in_attrs,
+ std::vector<int> *out_attrs) {
+ CHECK_EQ(in_attrs->size(), 3);
+
+ *dispatch_mode = DispatchMode::kFCompute;
+#if MXNET_USE_MKLDNN == 1
+ const PoolingParam ¶m = nnvm::get<PoolingParam>(attrs.parsed);
+ if (dev_mask == mshadow::cpu::kDevMask && SupportMKLDNNPooling(param)) {
+ *dispatch_mode = DispatchMode::kFComputeEx;
+ }
+#else
+ CHECK_EQ(out_attrs->size(), 3);
+#endif
+ for (size_t i = 0; i < out_attrs->size(); i++)
+ (*out_attrs)[i] = kDefaultStorage;
+ return true;
+}
+
NNVM_REGISTER_OP(_contrib_quantized_pooling)
.describe(R"code(Pooling operator for input and output data type of int8.
The input and output data comes with min and max thresholds for quantizing
@@ -101,7 +129,7 @@ the float32 data into int8.
This operator only supports `pool_type` of `avg` or `max`.)code"
ADD_FILELINE)
.set_num_inputs(3)
.set_num_outputs(3)
-.set_attr_parser(ParamParser<PoolingParam>)
+.set_attr_parser(PoolingParamParser)
.set_attr<nnvm::FListInputNames>("FListInputNames",
[](const NodeAttrs& attrs) {
return std::vector<std::string>{"data", "min_data", "max_data"};
@@ -112,6 +140,7 @@ the float32 data into int8.
})
.set_attr<nnvm::FInferShape>("FInferShape", QuantizedPoolingShape)
.set_attr<nnvm::FInferType>("FInferType", QuantizedPoolingType)
+.set_attr<FInferStorageType>("FInferStorageType", QuantizedPoolingStorageType)
.set_attr<FNeedRequantize>("FNeedRequantize",
[](const NodeAttrs& attrs) {
const PoolingParam& param = nnvm::get<PoolingParam>(attrs.parsed);
diff --git a/src/operator/quantization/requantize.cc
b/src/operator/quantization/requantize.cc
index 83ea37b835c..5ce0ff0b020 100644
--- a/src/operator/quantization/requantize.cc
+++ b/src/operator/quantization/requantize.cc
@@ -24,11 +24,31 @@
*/
#include "./requantize-inl.h"
#include "./quantize-inl.h"
+#if MXNET_USE_MKLDNN == 1
+#include "./mkldnn/mkldnn_requantize-inl.h"
+#endif
namespace mxnet {
namespace op {
DMLC_REGISTER_PARAMETER(RequantizeParam);
+bool RequantizeStorageType(const nnvm::NodeAttrs& attrs,
+ const int dev_mask,
+ DispatchMode* dispatch_mode,
+ std::vector<int> *in_attrs,
+ std::vector<int> *out_attrs) {
+ *dispatch_mode = DispatchMode::kFCompute;
+#if MXNET_USE_MKLDNN == 1
+ if (dev_mask == mshadow::cpu::kDevMask) {
+ *dispatch_mode = DispatchMode::kFComputeEx;
+ }
+#endif
+ (*out_attrs)[0] = kDefaultStorage;
+ (*out_attrs)[1] = kDefaultStorage;
+ (*out_attrs)[2] = kDefaultStorage;
+ return true;
+}
+
NNVM_REGISTER_OP(_contrib_requantize)
.describe(R"code(Given data that is quantized in int32 and the corresponding
thresholds,
requantize the data into int8 using min and max thresholds either calculated
at runtime
@@ -43,7 +63,12 @@ inference accuracy.
.set_num_outputs(3)
.set_attr<nnvm::FInferShape>("FInferShape", QuantizeShape)
.set_attr<nnvm::FInferType>("FInferType", RequantizeType)
+.set_attr<FInferStorageType>("FInferStorageType", RequantizeStorageType)
+#if MXNET_USE_MKLDNN == 1
+.set_attr<FComputeEx>("FComputeEx<cpu>", MKLDNNRequantizeForward)
+#else
.set_attr<FCompute>("FCompute<cpu>", RequantizeForward<cpu>)
+#endif
.set_attr<FResourceRequest>("FResourceRequest", [](const NodeAttrs& attrs) {
const RequantizeParam& param =
nnvm::get<RequantizeParam>(attrs.parsed);
diff --git a/tests/python/mkl/test_quantization_mkldnn.py
b/tests/python/mkl/test_quantization_mkldnn.py
new file mode 100644
index 00000000000..290f1a195c2
--- /dev/null
+++ b/tests/python/mkl/test_quantization_mkldnn.py
@@ -0,0 +1,28 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import os
+import sys
+import mxnet as mx
+
+os.environ['ENABLE_MKLDNN_QUANTIZATION_TEST'] = '1'
+curr_path = os.path.dirname(os.path.abspath(os.path.expanduser(__file__)))
+sys.path.insert(0, os.path.join(curr_path, '../quantization'))
+from test_quantization import *
+
+if __name__ == '__main__':
+ import nose
+ nose.runmodule()
diff --git a/tests/python/quantization/test_quantization.py
b/tests/python/quantization/test_quantization.py
index 7b08f46e836..15e8582b9ee 100644
--- a/tests/python/quantization/test_quantization.py
+++ b/tests/python/quantization/test_quantization.py
@@ -18,6 +18,7 @@
"""Some of the tests using CUDNN require a special GPU instruction called dp4a.
Ref:
http://images.nvidia.com/content/pdf/tesla/184457-Tesla-P4-Datasheet-NV-Final-Letter-Web.pdf
"""
+import os
import mxnet as mx
import numpy as np
from mxnet.test_utils import assert_almost_equal, rand_ndarray, rand_shape_nd,
same, DummyIter
@@ -25,6 +26,16 @@
from mxnet.module import Module
from mxnet.io import NDArrayIter
+def is_test_for_gpu():
+ return mx.current_context().device_type == 'gpu'
+
+def is_test_for_mkldnn():
+ return (mx.current_context().device_type == 'cpu'
+ and os.environ.get('ENABLE_MKLDNN_QUANTIZATION_TEST') == '1')
+
+def is_test_for_native_cpu():
+ return (mx.current_context().device_type == 'cpu'
+ and os.environ.get('ENABLE_MKLDNN_QUANTIZATION_TEST') == None)
@with_seed()
def test_quantize_float32_to_int8():
@@ -120,187 +131,220 @@ def check_requantize(shape, min_calib_range=None,
max_calib_range=None):
@with_seed()
def test_quantized_conv():
- if mx.current_context().device_type != 'gpu':
- print('skipped testing quantized_conv on cpu since it is not
implemented yet')
- return
-
- def check_quantized_conv(data_shape, kernel, num_filter, pad, stride,
no_bias):
- with mx.Context('gpu', 0):
- # run fp32 conv
- data = mx.sym.Variable(name='data', shape=data_shape,
dtype='float32')
- conv2d = mx.sym.Convolution(data=data, kernel=kernel,
num_filter=num_filter, pad=pad, stride=stride,
- no_bias=no_bias, cudnn_off=False,
name='conv2d')
- arg_shapes, _, _ = conv2d.infer_shape(data=data_shape)
- arg_names = conv2d.list_arguments()
- conv_exe_fp32 = conv2d.simple_bind(ctx=mx.current_context(),
grad_req='null')
- conv_exe_fp32.arg_dict[arg_names[0]][:] =
mx.nd.random.uniform(low=-127.0, high=127.0,
-
shape=data_shape).astype('int32')
- conv_exe_fp32.arg_dict[arg_names[1]][:] =
mx.nd.random.uniform(low=-127.0, high=127.0,
-
shape=arg_shapes[1]).astype('int32')
- if not no_bias:
- conv_exe_fp32.arg_dict[arg_names[2]][:] =
mx.nd.random.uniform(low=-127.0, high=127.0,
-
shape=arg_shapes[2]).astype('int32')
- output = conv_exe_fp32.forward()[0]
-
- # run quantized conv
- qdata = mx.sym.Variable(name='qdata', shape=data_shape,
dtype='int8')
- qweight = mx.sym.Variable(name='qweight', dtype='int8')
- min_data = mx.sym.Variable(name='min_data')
- max_data = mx.sym.Variable(name='max_data')
- min_weight = mx.sym.Variable(name='min_weight')
- max_weight = mx.sym.Variable(name='max_weight')
- quantized_conv2d = mx.sym.contrib.quantized_conv(data=qdata,
weight=qweight, min_data=min_data,
-
max_data=max_data, min_weight=min_weight,
-
max_weight=max_weight, kernel=kernel,
-
num_filter=num_filter, pad=pad, stride=stride,
- no_bias=no_bias)
- qarg_names = quantized_conv2d.list_arguments()
- type_dict = None
- if not no_bias:
- type_dict = {qarg_names[2]: 'int8'}
- conv_exe_int8 =
quantized_conv2d.simple_bind(ctx=mx.current_context(), type_dict=type_dict,
grad_req='null')
- conv_exe_int8.arg_dict[qarg_names[0]][:] =
conv_exe_fp32.arg_dict[arg_names[0]].astype('int8')
- conv_exe_int8.arg_dict[qarg_names[1]][:] =
conv_exe_fp32.arg_dict[arg_names[1]].astype('int8')
- quantized_range = 127.0
- if no_bias:
- conv_exe_int8.arg_dict[qarg_names[2]][:] = -quantized_range
- conv_exe_int8.arg_dict[qarg_names[3]][:] = quantized_range
- conv_exe_int8.arg_dict[qarg_names[4]][:] = -quantized_range
- conv_exe_int8.arg_dict[qarg_names[5]][:] = quantized_range
- else:
- conv_exe_int8.arg_dict[qarg_names[2]][:] =
conv_exe_fp32.arg_dict[arg_names[2]].astype('int8')
- conv_exe_int8.arg_dict[qarg_names[3]][:] = -quantized_range
- conv_exe_int8.arg_dict[qarg_names[4]][:] = quantized_range
- conv_exe_int8.arg_dict[qarg_names[5]][:] = -quantized_range
- conv_exe_int8.arg_dict[qarg_names[6]][:] = quantized_range
- conv_exe_int8.arg_dict[qarg_names[7]][:] = -quantized_range
- conv_exe_int8.arg_dict[qarg_names[8]][:] = quantized_range
- qoutput, min_range, max_range = conv_exe_int8.forward()
-
- if no_bias:
- assert_almost_equal(output.asnumpy(), qoutput.asnumpy())
- else:
- # with adding bias, accuracy loss should not be greater than
one
- diff = mx.nd.abs(output - qoutput.astype(output.dtype))
- cond = mx.nd.lesser(2, diff).sum().asscalar()
- assert cond == 0
-
- check_quantized_conv((3, 4, 28, 28), (3, 3), 128, (1, 1), (1, 1), True)
- check_quantized_conv((3, 4, 28, 28), (3, 3), 128, (1, 1), (1, 1), False)
+ def check_quantized_conv(data_shape, kernel, num_filter, pad, stride,
no_bias, qdtype):
+ if is_test_for_native_cpu():
+ print('skipped testing quantized_conv for native cpu since it is
not supported yet')
+ return
+ elif qdtype == 'int8' and is_test_for_mkldnn():
+ print('skipped testing quantized_conv for mkldnn cpu int8 since it
is not supported yet')
+ return
+ elif qdtype == 'uint8' and is_test_for_gpu():
+ print('skipped testing quantized_conv for gpu uint8 since it is
not supported yet')
+ return
+
+ # run fp32 conv
+ data = mx.sym.Variable(name='data', shape=data_shape, dtype='float32')
+ conv2d = mx.sym.Convolution(data=data, kernel=kernel,
num_filter=num_filter, pad=pad, stride=stride,
+ no_bias=no_bias, cudnn_off=False,
name='conv2d')
+ arg_shapes, _, _ = conv2d.infer_shape(data=data_shape)
+ arg_names = conv2d.list_arguments()
+ conv_exe_fp32 = conv2d.simple_bind(ctx=mx.current_context(),
grad_req='null')
+ if qdtype == 'uint8':
+ data_low = 0.0
+ data_high = 127.0
+ else:
+ data_low = -127.0
+ data_high = 127.0
+ conv_exe_fp32.arg_dict[arg_names[0]][:] =
mx.nd.random.uniform(low=data_low, high=data_high,
+
shape=data_shape).astype('int32')
+ conv_exe_fp32.arg_dict[arg_names[1]][:] =
mx.nd.random.uniform(low=-127.0, high=127.0,
+
shape=arg_shapes[1]).astype('int32')
+ if not no_bias:
+ conv_exe_fp32.arg_dict[arg_names[2]][:] =
mx.nd.random.uniform(low=-127.0, high=127.0,
+
shape=arg_shapes[2]).astype('int32')
+ output = conv_exe_fp32.forward()[0]
+
+ # run quantized conv
+ qdata = mx.sym.Variable(name='qdata', shape=data_shape, dtype=qdtype)
+ qweight = mx.sym.Variable(name='qweight', dtype='int8')
+ min_data = mx.sym.Variable(name='min_data')
+ max_data = mx.sym.Variable(name='max_data')
+ min_weight = mx.sym.Variable(name='min_weight')
+ max_weight = mx.sym.Variable(name='max_weight')
+ quantized_conv2d = mx.sym.contrib.quantized_conv(data=qdata,
weight=qweight, min_data=min_data,
+ max_data=max_data,
min_weight=min_weight,
+
max_weight=max_weight, kernel=kernel,
+
num_filter=num_filter, pad=pad, stride=stride,
+ no_bias=no_bias)
+ qarg_names = quantized_conv2d.list_arguments()
+ type_dict = None
+ if not no_bias:
+ type_dict = {qarg_names[2]: 'int8'}
+ conv_exe_int8 = quantized_conv2d.simple_bind(ctx=mx.current_context(),
type_dict=type_dict, grad_req='null')
+ conv_exe_int8.arg_dict[qarg_names[0]][:] =
conv_exe_fp32.arg_dict[arg_names[0]].astype(qdtype)
+ conv_exe_int8.arg_dict[qarg_names[1]][:] =
conv_exe_fp32.arg_dict[arg_names[1]].astype('int8')
+ quantized_range = 127.0
+ if no_bias:
+ conv_exe_int8.arg_dict[qarg_names[2]][:] = -quantized_range
+ conv_exe_int8.arg_dict[qarg_names[3]][:] = quantized_range
+ conv_exe_int8.arg_dict[qarg_names[4]][:] = -quantized_range
+ conv_exe_int8.arg_dict[qarg_names[5]][:] = quantized_range
+ else:
+ conv_exe_int8.arg_dict[qarg_names[2]][:] =
conv_exe_fp32.arg_dict[arg_names[2]].astype('int8')
+ conv_exe_int8.arg_dict[qarg_names[3]][:] = -quantized_range
+ conv_exe_int8.arg_dict[qarg_names[4]][:] = quantized_range
+ conv_exe_int8.arg_dict[qarg_names[5]][:] = -quantized_range
+ conv_exe_int8.arg_dict[qarg_names[6]][:] = quantized_range
+ conv_exe_int8.arg_dict[qarg_names[7]][:] = -quantized_range
+ conv_exe_int8.arg_dict[qarg_names[8]][:] = quantized_range
+ qoutput, min_range, max_range = conv_exe_int8.forward()
+
+ if no_bias:
+ assert_almost_equal(output.asnumpy(), qoutput.asnumpy())
+ else:
+ # with adding bias, accuracy loss should not be greater than one
+ diff = mx.nd.abs(output - qoutput.astype(output.dtype))
+ cond = mx.nd.lesser(2, diff).sum().asscalar()
+ assert cond == 0
+ for qdtype in ['int8', 'uint8']:
+ check_quantized_conv((3, 4, 28, 28), (3, 3), 128, (1, 1), (1, 1),
True, qdtype)
+ check_quantized_conv((3, 4, 28, 28), (3, 3), 128, (1, 1), (1, 1),
False, qdtype)
@with_seed()
def test_quantized_pooling():
- if mx.current_context().device_type != 'gpu':
- print('skipped testing quantized_pooling on cpu since it is not
implemented yet')
- return
-
- def check_quantized_pooling(data_shape, kernel, pool_type, pad, stride,
global_pool):
- with mx.Context('gpu', 0):
- data = mx.sym.Variable(name='data', shape=data_shape,
dtype='float32')
- pooling_fp32 = mx.sym.Pooling(data=data, kernel=kernel, pad=pad,
stride=stride,
- pool_type=pool_type,
global_pool=global_pool, cudnn_off=False)
- arg_shapes, _, _ = pooling_fp32.infer_shape(data=data_shape)
- arg_names = pooling_fp32.list_arguments()
- pooling_fp32_exe =
pooling_fp32.simple_bind(ctx=mx.current_context(), grad_req='null')
- pooling_fp32_exe.arg_dict[arg_names[0]][:] =
mx.nd.random.uniform(low=-127.0, high=127.0,
-
shape=data_shape).astype('int32')
- output = pooling_fp32_exe.forward()[0]
-
- qdata = mx.sym.Variable(name='qdata', shape=data_shape,
dtype='int8')
- min_data = mx.sym.Variable(name='min_data')
- max_data = mx.sym.Variable(name='max_data')
- quantized_pooling = mx.sym.contrib.quantized_pooling(data=qdata,
min_data=min_data,
-
max_data=max_data, kernel=kernel,
- pad=pad,
stride=stride, pool_type=pool_type,
-
global_pool=global_pool)
- pooling_int8_exe =
quantized_pooling.simple_bind(ctx=mx.current_context(), grad_req='null')
- qarg_names = quantized_pooling.list_arguments()
- pooling_int8_exe.arg_dict[qarg_names[0]][:] =
pooling_fp32_exe.arg_dict[arg_names[0]].astype('int8')
- quantized_range = 127.0
- pooling_int8_exe.arg_dict[qarg_names[1]][:] = -quantized_range
- pooling_int8_exe.arg_dict[qarg_names[2]][:] = quantized_range
- qoutput, min_range, max_range = pooling_int8_exe.forward()
-
- if pool_type == 'max':
- assert_almost_equal(output.asnumpy(), qoutput.asnumpy())
- elif pool_type == 'avg': # for avg pooling, fp32 and int8 may be
different due to rounding errors
- diff = mx.nd.abs(output - qoutput.astype(output.dtype))
- cond = mx.nd.lesser(2, diff).sum().asscalar()
- assert cond == 0
-
- check_quantized_pooling((3, 4, 56, 56), (3, 3), 'max', (0, 0), (2, 2),
False)
- check_quantized_pooling((3, 4, 56, 56), (3, 3), 'max', (0, 0), (2, 2),
True)
- check_quantized_pooling((3, 512, 7, 7), (7, 7), 'avg', (0, 0), (1, 1),
False)
- check_quantized_pooling((3, 512, 7, 7), (7, 7), 'avg', (0, 0), (1, 1),
True)
-
+ def check_quantized_pooling(data_shape, kernel, pool_type, pad, stride,
global_pool, qdtype):
+ if is_test_for_native_cpu():
+ print('skipped testing quantized_pooling for native cpu since it
is not supported yet')
+ return
+ elif qdtype == 'uint8' and is_test_for_gpu():
+ print('skipped testing quantized_pooling for gpu uint8 since it is
not supported yet')
+ return
+
+ data = mx.sym.Variable(name='data', shape=data_shape, dtype='float32')
+ pooling_fp32 = mx.sym.Pooling(data=data, kernel=kernel, pad=pad,
stride=stride,
+ pool_type=pool_type,
global_pool=global_pool, cudnn_off=False)
+ arg_shapes, _, _ = pooling_fp32.infer_shape(data=data_shape)
+ arg_names = pooling_fp32.list_arguments()
+ pooling_fp32_exe = pooling_fp32.simple_bind(ctx=mx.current_context(),
grad_req='null')
+ if qdtype == 'uint8':
+ data_low = 0.0
+ data_high = 127.0
+ else:
+ data_low = -127.0
+ data_high = 127.0
+ pooling_fp32_exe.arg_dict[arg_names[0]][:] =
mx.nd.random.uniform(low=data_low, high=data_high,
+
shape=data_shape).astype('int32')
+ output = pooling_fp32_exe.forward()[0]
+
+ qdata = mx.sym.Variable(name='qdata', shape=data_shape, dtype=qdtype)
+ min_data = mx.sym.Variable(name='min_data')
+ max_data = mx.sym.Variable(name='max_data')
+ quantized_pooling = mx.sym.contrib.quantized_pooling(data=qdata,
min_data=min_data,
+
max_data=max_data, kernel=kernel,
+ pad=pad,
stride=stride, pool_type=pool_type,
+
global_pool=global_pool)
+ pooling_int8_exe =
quantized_pooling.simple_bind(ctx=mx.current_context(), grad_req='null')
+ qarg_names = quantized_pooling.list_arguments()
+ pooling_int8_exe.arg_dict[qarg_names[0]][:] =
pooling_fp32_exe.arg_dict[arg_names[0]].astype(qdtype)
+ quantized_range = 127.0
+ pooling_int8_exe.arg_dict[qarg_names[1]][:] = -quantized_range
+ pooling_int8_exe.arg_dict[qarg_names[2]][:] = quantized_range
+ qoutput, min_range, max_range = pooling_int8_exe.forward()
+
+ if pool_type == 'max':
+ assert_almost_equal(output.asnumpy(), qoutput.asnumpy())
+ elif pool_type == 'avg': # for avg pooling, fp32 and int8 may be
different due to rounding errors
+ diff = mx.nd.abs(output - qoutput.astype(output.dtype))
+ cond = mx.nd.lesser(2, diff).sum().asscalar()
+ assert cond == 0
+
+ for qdtype in ['int8', 'uint8']:
+ check_quantized_pooling((3, 4, 56, 56), (3, 3), 'max', (0, 0), (2, 2),
False, qdtype)
+ check_quantized_pooling((3, 4, 56, 56), (3, 3), 'max', (0, 0), (2, 2),
True, qdtype)
+ check_quantized_pooling((3, 512, 7, 7), (7, 7), 'avg', (0, 0), (1, 1),
False, qdtype)
+ check_quantized_pooling((3, 512, 7, 7), (7, 7), 'avg', (0, 0), (1, 1),
True, qdtype)
@with_seed()
def test_quantized_fc():
- if mx.current_context().device_type != 'gpu':
- print('skipped testing quantized_fc on cpu since it is not implemented
yet')
- return
-
- def check_quantized_fc(data_shape, num_hidden, no_bias, flatten=True):
- with mx.Context('gpu', 0):
- data = mx.sym.Variable(name='data', shape=data_shape,
dtype='float32')
- fc_fp32 = mx.sym.FullyConnected(data=data, num_hidden=num_hidden,
no_bias=no_bias, flatten=flatten)
- arg_shapes, _, _ = fc_fp32.infer_shape(data=data_shape)
- arg_names = fc_fp32.list_arguments()
- fc_fp32_exe = fc_fp32.simple_bind(ctx=mx.current_context(),
grad_req='null')
- fc_fp32_exe.arg_dict[arg_names[0]][:] =
mx.nd.random.uniform(low=-127.0, high=127.0,
-
shape=data_shape).astype('int32')
- fc_fp32_exe.arg_dict[arg_names[1]][:] =
mx.nd.random.uniform(low=-127.0, high=127.0,
-
shape=arg_shapes[1]).astype('int32')
- if not no_bias:
- fc_fp32_exe.arg_dict[arg_names[2]][:] =
mx.nd.random.uniform(low=-127.0, high=127.0,
-
shape=arg_shapes[2]).astype('int32')
- output = fc_fp32_exe.forward()[0]
-
- qdata = mx.sym.Variable(name='qdata', shape=data_shape,
dtype='int8')
- fc_int8 = mx.sym.contrib.quantized_fully_connected(data=qdata,
num_hidden=num_hidden,
-
no_bias=no_bias, flatten=flatten)
- qarg_names = fc_int8.list_arguments()
- type_dict = {qarg_names[1]: 'int8'}
- if not no_bias:
- type_dict.update({qarg_names[2]: 'int8'})
- fc_int8_exe = fc_int8.simple_bind(ctx=mx.current_context(),
type_dict=type_dict, grad_req='null')
- fc_int8_exe.arg_dict[qarg_names[0]][:] =
fc_fp32_exe.arg_dict[arg_names[0]].astype('int8')
- fc_int8_exe.arg_dict[qarg_names[1]][:] =
fc_fp32_exe.arg_dict[arg_names[1]].astype('int8')
- quantized_range = 127.0
- if no_bias:
- fc_int8_exe.arg_dict[qarg_names[2]][:] = -quantized_range
- fc_int8_exe.arg_dict[qarg_names[3]][:] = quantized_range
- fc_int8_exe.arg_dict[qarg_names[4]][:] = -quantized_range
- fc_int8_exe.arg_dict[qarg_names[5]][:] = quantized_range
- else:
- fc_int8_exe.arg_dict[qarg_names[2]][:] =
fc_fp32_exe.arg_dict[arg_names[2]].astype('int8')
- fc_int8_exe.arg_dict[qarg_names[3]][:] = -quantized_range
- fc_int8_exe.arg_dict[qarg_names[4]][:] = quantized_range
- fc_int8_exe.arg_dict[qarg_names[5]][:] = -quantized_range
- fc_int8_exe.arg_dict[qarg_names[6]][:] = quantized_range
- fc_int8_exe.arg_dict[qarg_names[7]][:] = -quantized_range
- fc_int8_exe.arg_dict[qarg_names[8]][:] = quantized_range
- qoutput, min_range, max_range = fc_int8_exe.forward()
-
- if no_bias:
- assert_almost_equal(output.asnumpy(), qoutput.asnumpy())
- else:
- # with adding bias, accuracy loss should not be greater than
one
- diff = mx.nd.abs(output - qoutput.astype(output.dtype))
- cond = mx.nd.lesser(2, diff).sum().asscalar()
- assert cond == 0
-
- check_quantized_fc((32, 512, 2, 2), 100, True)
- check_quantized_fc((32, 111, 2, 2), 100, True)
- check_quantized_fc((32, 512, 2, 2), 100, False)
- check_quantized_fc((32, 111, 2, 2), 100, False)
+ def check_quantized_fc(data_shape, num_hidden, no_bias, qdtype,
flatten=True):
+ if mx.current_context().device_type != 'gpu':
+ print('skipped testing quantized_fc on cpu since it is not
supported yet')
+ return
+ elif qdtype == 'uint8' and is_test_for_gpu():
+ print('skipped testing quantized_fc for gpu uint8 since it is not
supported yet')
+ return
+
+ data = mx.sym.Variable(name='data', shape=data_shape, dtype='float32')
+ fc_fp32 = mx.sym.FullyConnected(data=data, num_hidden=num_hidden,
no_bias=no_bias, flatten=flatten)
+ arg_shapes, _, _ = fc_fp32.infer_shape(data=data_shape)
+ arg_names = fc_fp32.list_arguments()
+ fc_fp32_exe = fc_fp32.simple_bind(ctx=mx.current_context(),
grad_req='null')
+ if qdtype == 'uint8':
+ data_low = 0.0
+ data_high = 127.0
+ else:
+ data_low = -127.0
+ data_high = 127.0
+ fc_fp32_exe.arg_dict[arg_names[0]][:] =
mx.nd.random.uniform(low=data_low, high=data_high,
+
shape=data_shape).astype('int32')
+ fc_fp32_exe.arg_dict[arg_names[1]][:] =
mx.nd.random.uniform(low=-127.0, high=127.0,
+
shape=arg_shapes[1]).astype('int32')
+ if not no_bias:
+ fc_fp32_exe.arg_dict[arg_names[2]][:] =
mx.nd.random.uniform(low=-127.0, high=127.0,
+
shape=arg_shapes[2]).astype('int32')
+ output = fc_fp32_exe.forward()[0]
+
+ qdata = mx.sym.Variable(name='qdata', shape=data_shape, dtype='int8')
+ fc_int8 = mx.sym.contrib.quantized_fully_connected(data=qdata,
num_hidden=num_hidden,
+ no_bias=no_bias,
flatten=flatten)
+ qarg_names = fc_int8.list_arguments()
+ type_dict = {qarg_names[1]: 'int8'}
+ if not no_bias:
+ type_dict.update({qarg_names[2]: 'int8'})
+ fc_int8_exe = fc_int8.simple_bind(ctx=mx.current_context(),
type_dict=type_dict, grad_req='null')
+ fc_int8_exe.arg_dict[qarg_names[0]][:] =
fc_fp32_exe.arg_dict[arg_names[0]].astype(qdtype)
+ fc_int8_exe.arg_dict[qarg_names[1]][:] =
fc_fp32_exe.arg_dict[arg_names[1]].astype('int8')
+ quantized_range = 127.0
+ if no_bias:
+ fc_int8_exe.arg_dict[qarg_names[2]][:] = -quantized_range
+ fc_int8_exe.arg_dict[qarg_names[3]][:] = quantized_range
+ fc_int8_exe.arg_dict[qarg_names[4]][:] = -quantized_range
+ fc_int8_exe.arg_dict[qarg_names[5]][:] = quantized_range
+ else:
+ fc_int8_exe.arg_dict[qarg_names[2]][:] =
fc_fp32_exe.arg_dict[arg_names[2]].astype('int8')
+ fc_int8_exe.arg_dict[qarg_names[3]][:] = -quantized_range
+ fc_int8_exe.arg_dict[qarg_names[4]][:] = quantized_range
+ fc_int8_exe.arg_dict[qarg_names[5]][:] = -quantized_range
+ fc_int8_exe.arg_dict[qarg_names[6]][:] = quantized_range
+ fc_int8_exe.arg_dict[qarg_names[7]][:] = -quantized_range
+ fc_int8_exe.arg_dict[qarg_names[8]][:] = quantized_range
+ qoutput, min_range, max_range = fc_int8_exe.forward()
+
+ if no_bias:
+ assert_almost_equal(output.asnumpy(), qoutput.asnumpy())
+ else:
+ # with adding bias, accuracy loss should not be greater than one
+ diff = mx.nd.abs(output - qoutput.astype(output.dtype))
+ cond = mx.nd.lesser(2, diff).sum().asscalar()
+ assert cond == 0
+ for qdtype in ['int8', 'uint8']:
+ check_quantized_fc((32, 512, 2, 2), 100, True, qdtype)
+ check_quantized_fc((32, 111, 2, 2), 100, True, qdtype)
+ check_quantized_fc((32, 512, 2, 2), 100, False, qdtype)
+ check_quantized_fc((32, 111, 2, 2), 100, False, qdtype)
@with_seed()
def test_quantized_flatten():
- def check_quantized_flatten(shape):
- qdata = mx.nd.random.uniform(low=-127, high=127,
shape=shape).astype('int8')
+ def check_quantized_flatten(shape, qdtype):
+ if qdtype == 'uint8':
+ data_low = 0.0
+ data_high = 127.0
+ else:
+ data_low = -127.0
+ data_high = 127.0
+ qdata = mx.nd.random.uniform(low=data_low, high=data_high,
shape=shape).astype(qdtype)
min_data = mx.nd.array([-1023.343], dtype='float32')
max_data = mx.nd.array([2343.324275], dtype='float32')
qoutput, min_output, max_output =
mx.nd.contrib.quantized_flatten(qdata, min_data, max_data)
@@ -311,10 +355,11 @@ def check_quantized_flatten(shape):
assert same(min_data.asnumpy(), min_output.asnumpy())
assert same(max_data.asnumpy(), max_output.asnumpy())
- check_quantized_flatten((10,))
- check_quantized_flatten((10, 15))
- check_quantized_flatten((10, 15, 18))
- check_quantized_flatten((3, 4, 23, 23))
+ for qdtype in ['int8', 'uint8']:
+ check_quantized_flatten((10,), qdtype)
+ check_quantized_flatten((10, 15), qdtype)
+ check_quantized_flatten((10, 15, 18), qdtype)
+ check_quantized_flatten((3, 4, 23, 23), qdtype)
@with_seed()
@@ -353,56 +398,69 @@ def get_fp32_sym():
@with_seed()
def test_quantize_model():
- def check_params(params, qparams, qsym=None):
- if qsym is None:
- assert len(params) == len(qparams)
- for k, v in params.items():
- assert k in qparams
- assert same(v.asnumpy(), qparams[k].asnumpy())
- else:
- qparams_ground_truth = mx.contrib.quant._quantize_params(qsym,
params)
- assert len(qparams) == len(qparams_ground_truth)
- for k, v in qparams_ground_truth.items():
- assert k in qparams
- assert same(v.asnumpy(), qparams[k].asnumpy())
-
- def check_qsym_calibrated(qsym):
- attrs = qsym.attr_dict()
- for k, v in attrs.items():
- if k.find('requantize_') != -1:
- assert 'min_calib_range' in v
- assert 'max_calib_range' in v
-
- sym = get_fp32_sym()
- mod = Module(symbol=sym)
- batch_size = 4
- data_shape = (batch_size, 4, 10, 10)
- label_shape = (batch_size, 10)
- mod.bind(data_shapes=[('data', data_shape)],
label_shapes=[('softmax_label', label_shape)])
- mod.init_params()
- arg_params, aux_params = mod.get_params()
- qsym, qarg_params, qaux_params = mx.contrib.quant.quantize_model(sym=sym,
-
arg_params=arg_params,
-
aux_params=aux_params,
-
ctx=mx.current_context(),
-
calib_mode='none')
- check_params(arg_params, qarg_params, qsym)
- check_params(aux_params, qaux_params)
-
- calib_data = mx.nd.random.uniform(shape=data_shape)
- calib_data = NDArrayIter(data=calib_data)
- calib_data = DummyIter(calib_data)
- qsym, qarg_params, qaux_params = mx.contrib.quant.quantize_model(sym=sym,
-
arg_params=arg_params,
-
aux_params=aux_params,
-
ctx=mx.current_context(),
-
calib_mode='naive',
-
calib_data=calib_data,
-
num_calib_examples=20)
- check_params(arg_params, qarg_params, qsym)
- check_params(aux_params, qaux_params)
- check_qsym_calibrated(qsym)
-
+ def check_quantize_model(qdtype):
+ def check_params(params, qparams, qsym=None):
+ if qsym is None:
+ assert len(params) == len(qparams)
+ for k, v in params.items():
+ assert k in qparams
+ assert same(v.asnumpy(), qparams[k].asnumpy())
+ else:
+ qparams_ground_truth = mx.contrib.quant._quantize_params(qsym,
params)
+ assert len(qparams) == len(qparams_ground_truth)
+ for k, v in qparams_ground_truth.items():
+ assert k in qparams
+ assert same(v.asnumpy(), qparams[k].asnumpy())
+
+ def check_qsym_calibrated(qsym):
+ attrs = qsym.attr_dict()
+ for k, v in attrs.items():
+ if k.find('requantize_') != -1:
+ assert 'min_calib_range' in v
+ assert 'max_calib_range' in v
+
+ def check_qsym_qdtype(qsym, qdtype):
+ attrs = qsym.attr_dict()
+ for k, v in attrs.items():
+ if k.find('_quantize') != -1:
+ assert 'out_type' in v
+ assert v['out_type'] == qdtype
+
+ sym = get_fp32_sym()
+ mod = Module(symbol=sym)
+ batch_size = 4
+ data_shape = (batch_size, 4, 10, 10)
+ label_shape = (batch_size, 10)
+ mod.bind(data_shapes=[('data', data_shape)],
label_shapes=[('softmax_label', label_shape)])
+ mod.init_params()
+ arg_params, aux_params = mod.get_params()
+ qsym, qarg_params, qaux_params =
mx.contrib.quant.quantize_model(sym=sym,
+
arg_params=arg_params,
+
aux_params=aux_params,
+
ctx=mx.current_context(),
+
quantized_dtype=qdtype,
+
calib_mode='none')
+ check_params(arg_params, qarg_params, qsym)
+ check_params(aux_params, qaux_params)
+
+ calib_data = mx.nd.random.uniform(shape=data_shape)
+ calib_data = NDArrayIter(data=calib_data)
+ calib_data = DummyIter(calib_data)
+ qsym, qarg_params, qaux_params =
mx.contrib.quant.quantize_model(sym=sym,
+
arg_params=arg_params,
+
aux_params=aux_params,
+
ctx=mx.current_context(),
+
quantized_dtype=qdtype,
+
calib_mode='naive',
+
calib_data=calib_data,
+
num_calib_examples=20)
+ check_params(arg_params, qarg_params, qsym)
+ check_params(aux_params, qaux_params)
+ check_qsym_calibrated(qsym)
+ check_qsym_qdtype(qsym, qdtype)
+
+ for qdtype in ['int8', 'uint8']:
+ check_quantize_model(qdtype)
@with_seed()
def test_quantize_sym_with_calib():
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
