From 3e574eeb6e1cf24ce2456ba0c7a7c90b5f01020c Mon Sep 17 00:00:00 2001 From: xinyu-intel Date: Tue, 19 Jun 2018 15:14:59 +0800 Subject: [PATCH 1/4] add softmax imporvement --- src/operator/nn/mkldnn/mkldnn_softmax.cc | 8 +++++++- src/operator/nn/softmax.cc | 3 +-- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/src/operator/nn/mkldnn/mkldnn_softmax.cc b/src/operator/nn/mkldnn/mkldnn_softmax.cc index aa59f13d06da..6f31f789f9a7 100644 --- a/src/operator/nn/mkldnn/mkldnn_softmax.cc +++ b/src/operator/nn/mkldnn/mkldnn_softmax.cc @@ -38,11 +38,17 @@ void MKLDNNSoftmaxForward(const nnvm::NodeAttrs& attrs, const OpContext &ctx, auto input_mem = in_data.GetMKLDNNData(); mkldnn::memory::primitive_desc data_mpd = input_mem->get_primitive_desc(); mkldnn::memory::desc data_md = data_mpd.desc(); + int axis = param.axis; + int ndim = in_data.shape().ndim(); + CHECK(axis < ndim && axis >= -ndim) + << "axis " << axis << "exceeds the input dimension of " << ndim; + axis = (axis + ndim) % ndim; + auto cpu_engine = data_mpd.get_engine(); auto prop = ctx.is_train ? mkldnn::prop_kind::forward_training : mkldnn::prop_kind::forward_scoring; mkldnn::softmax_forward::desc desc = mkldnn::softmax_forward::desc(prop, - data_md, param.axis); + data_md, axis); mkldnn::softmax_forward::primitive_desc pdesc(desc, cpu_engine); auto output_memory = out_data.GetMKLDNNData(); diff --git a/src/operator/nn/softmax.cc b/src/operator/nn/softmax.cc index f8cc6fee9a22..b0742820369e 100644 --- a/src/operator/nn/softmax.cc +++ b/src/operator/nn/softmax.cc @@ -38,10 +38,9 @@ static void SoftmaxComputeExCPU(const nnvm::NodeAttrs& attrs, const std::vector& inputs, const std::vector& req, const std::vector& outputs) { - const SoftmaxParam& param = nnvm::get(attrs.parsed); // It seems MKLDNN softmax doesn't support training. // and it only supports non-negative axis. - if (SupportMKLDNN(inputs[0]) && !ctx.is_train && param.axis >= 0) { + if (SupportMKLDNN(inputs[0]) && !ctx.is_train) { MKLDNN_OPCHECK_INIT(false, outputs.size(), inputs, outputs); MKLDNNSoftmaxForward(attrs, ctx, inputs[0], req[0], outputs[0]); auto fn = SoftmaxCompute; From 933667d7ee403873e7e7bf81ec820ae728cc87b4 Mon Sep 17 00:00:00 2001 From: xinyu-intel Date: Tue, 19 Jun 2018 17:23:33 +0800 Subject: [PATCH 2/4] reuse CheckAxis code --- src/operator/nn/mkldnn/mkldnn_softmax.cc | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/src/operator/nn/mkldnn/mkldnn_softmax.cc b/src/operator/nn/mkldnn/mkldnn_softmax.cc index 6f31f789f9a7..acfa358a796e 100644 --- a/src/operator/nn/mkldnn/mkldnn_softmax.cc +++ b/src/operator/nn/mkldnn/mkldnn_softmax.cc @@ -26,6 +26,7 @@ #include "../softmax-inl.h" #include "./mkldnn_ops-inl.h" #include "./mkldnn_base-inl.h" +#include "../../tensor/broadcast_reduce_op.h" #if MXNET_USE_MKLDNN == 1 namespace mxnet { @@ -38,11 +39,7 @@ void MKLDNNSoftmaxForward(const nnvm::NodeAttrs& attrs, const OpContext &ctx, auto input_mem = in_data.GetMKLDNNData(); mkldnn::memory::primitive_desc data_mpd = input_mem->get_primitive_desc(); mkldnn::memory::desc data_md = data_mpd.desc(); - int axis = param.axis; - int ndim = in_data.shape().ndim(); - CHECK(axis < ndim && axis >= -ndim) - << "axis " << axis << "exceeds the input dimension of " << ndim; - axis = (axis + ndim) % ndim; + int axis = CheckAxis(param.axis, in_data.shape().ndim()); auto cpu_engine = data_mpd.get_engine(); auto prop = ctx.is_train From 0580f86eaa451afa31bbae9900b5f6e8f9d0b2f5 Mon Sep 17 00:00:00 2001 From: xinyu-intel Date: Tue, 19 Jun 2018 18:34:30 +0800 Subject: [PATCH 3/4] update comment --- src/operator/nn/softmax.cc | 1 - 1 file changed, 1 deletion(-) diff --git a/src/operator/nn/softmax.cc b/src/operator/nn/softmax.cc index b0742820369e..e9b104f12868 100644 --- a/src/operator/nn/softmax.cc +++ b/src/operator/nn/softmax.cc @@ -39,7 +39,6 @@ static void SoftmaxComputeExCPU(const nnvm::NodeAttrs& attrs, const std::vector& req, const std::vector& outputs) { // It seems MKLDNN softmax doesn't support training. - // and it only supports non-negative axis. if (SupportMKLDNN(inputs[0]) && !ctx.is_train) { MKLDNN_OPCHECK_INIT(false, outputs.size(), inputs, outputs); MKLDNNSoftmaxForward(attrs, ctx, inputs[0], req[0], outputs[0]); From 02219ea9ccdb30cf86233f3120355bc2acaaaead Mon Sep 17 00:00:00 2001 From: xinyu-intel Date: Wed, 20 Jun 2018 09:57:50 +0800 Subject: [PATCH 4/4] add tests with negative axis --- tests/python/unittest/test_operator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/python/unittest/test_operator.py b/tests/python/unittest/test_operator.py index f287c1919636..674266934363 100644 --- a/tests/python/unittest/test_operator.py +++ b/tests/python/unittest/test_operator.py @@ -4098,7 +4098,7 @@ def test_new_softmax(): for ndim in range(1, 5): for _ in range(5): shape = np.random.randint(1, 5, size=ndim) - axis = np.random.randint(0, ndim) + axis = np.random.randint(-ndim, ndim) data = np.random.uniform(-2, 2, size=shape) sym = mx.sym.softmax(axis=axis) check_symbolic_forward(sym, [data], [np_softmax(data, axis=axis)])