This is an automated email from the ASF dual-hosted git repository.
zhasheng pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-mxnet.git
The following commit(s) were added to refs/heads/master by this push:
new 6307c00 Fix axis Bug in MKLDNN Softmax (#11335)
6307c00 is described below
commit 6307c00b1a9648e86d357d259d2068f53cc0a257
Author: Xinyu Chen <xinyu1.c...@intel.com>
AuthorDate: Wed Jun 20 12:09:03 2018 +0800
Fix axis Bug in MKLDNN Softmax (#11335)
* add softmax imporvement
* reuse CheckAxis code
* update comment
* add tests with negative axis
---
src/operator/nn/mkldnn/mkldnn_softmax.cc | 5 ++++-
src/operator/nn/softmax.cc | 4 +---
tests/python/unittest/test_operator.py | 2 +-
3 files changed, 6 insertions(+), 5 deletions(-)
diff --git a/src/operator/nn/mkldnn/mkldnn_softmax.cc
b/src/operator/nn/mkldnn/mkldnn_softmax.cc
index aa59f13..acfa358 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,13 @@ 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 = CheckAxis(param.axis, in_data.shape().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 f8cc6fe..e9b104f 100644
--- a/src/operator/nn/softmax.cc
+++ b/src/operator/nn/softmax.cc
@@ -38,10 +38,8 @@ static void SoftmaxComputeExCPU(const nnvm::NodeAttrs& attrs,
const std::vector<NDArray>& inputs,
const std::vector<OpReqType>& req,
const std::vector<NDArray>& outputs) {
- const SoftmaxParam& param = nnvm::get<SoftmaxParam>(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<cpu, mxnet_op::softmax_fwd>;
diff --git a/tests/python/unittest/test_operator.py
b/tests/python/unittest/test_operator.py
index f287c19..6742669 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)])
