diff --git a/src/operator/nn/activation-inl.h b/src/operator/nn/activation-inl.h index 89a369c6717e..32a7a5ad6176 100644 --- a/src/operator/nn/activation-inl.h +++ b/src/operator/nn/activation-inl.h @@ -201,7 +201,7 @@ void ActivationGradCompute(const nnvm::NodeAttrs& attrs, const std::vector& inputs, const std::vector& req, const std::vector& outputs) { -#if MXNET_USE_CUDNN == 1 +#if (MXNET_USE_CUDNN == 1 || MXNET_USE_MKLDNN == 1) CHECK_EQ(inputs.size(), 3U); #else CHECK_EQ(inputs.size(), 2U); diff --git a/src/operator/nn/activation.cc b/src/operator/nn/activation.cc index 89059321b693..08028265c485 100644 --- a/src/operator/nn/activation.cc +++ b/src/operator/nn/activation.cc @@ -44,7 +44,7 @@ struct ActivationGrad { const std::vector& ograds) const { std::vector heads(ograds.begin(), ograds.end()); heads.emplace_back(nnvm::NodeEntry{n, activation::kOut, 0}); -#if MXNET_USE_CUDNN == 1 +#if (MXNET_USE_CUDNN == 1 || MXNET_USE_MKLDNN == 1) heads.push_back(n->inputs[activation::kData]); #endif return MakeGradNode(op_name, n, heads, n->attrs.dict); @@ -74,15 +74,11 @@ void ActivationGradComputeExCPU(const nnvm::NodeAttrs& attrs, const std::vector& inputs, const std::vector& req, const std::vector& outputs) { -#if MXNET_USE_CUDNN == 1 CHECK_EQ(inputs.size(), 3U); -#else - CHECK_EQ(inputs.size(), 2U); -#endif const ActivationParam& param = nnvm::get(attrs.parsed); if (SupportMKLDNN(inputs[0])) { MKLDNN_OPCHECK_INIT(true, outputs.size(), inputs, outputs); - MKLDNNActivationBackward(attrs, ctx, inputs[0], inputs[1], req[0], + MKLDNNActivationBackward(attrs, ctx, inputs[0], inputs[2], req[0], outputs[0]); MKLDNN_OPCHECK_RUN(ActivationGradCompute, attrs, ctx, inputs, req, outputs); return; @@ -116,13 +112,13 @@ inline static bool BackwardActStorageType(const nnvm::NodeAttrs& attrs, DispatchMode* dispatch_mode, std::vector *in_attrs, std::vector *out_attrs) { -#if MXNET_USE_CUDNN == 1 +#if (MXNET_USE_CUDNN == 1 || MXNET_USE_MKLDNN == 1) CHECK_EQ(in_attrs->size(), 3U); #else CHECK_EQ(in_attrs->size(), 2U); #endif CHECK_EQ(out_attrs->size(), 1U); -#if MXNET_USE_CUDNN == 1 +#if (MXNET_USE_CUDNN == 1 || MXNET_USE_MKLDNN == 1) bool ret = ElemwiseStorageType<3, 1, false, false, false>(attrs, dev_mask, dispatch_mode, in_attrs, out_attrs); diff --git a/src/operator/nn/mkldnn/mkldnn_act.cc b/src/operator/nn/mkldnn/mkldnn_act.cc index 8c19850ced38..9be5bfbc150d 100644 --- a/src/operator/nn/mkldnn/mkldnn_act.cc +++ b/src/operator/nn/mkldnn/mkldnn_act.cc @@ -43,13 +43,10 @@ namespace mxnet { namespace op { bool SupportMKLDNNAct(const ActivationParam& param) { - // We only enable ReLU for now. It seems other activations have some precision - // problems. - return param.act_type == activation::kReLU; -#if 0 + return param.act_type == activation::kReLU || param.act_type == activation::kSigmoid - || param.act_type == activation::kSoftReLU; -#endif + || param.act_type == activation::kSoftReLU + || param.act_type == activation::kTanh; } static inline mkldnn::algorithm GetMKLDNNActAlgo(const ActivationParam& param) { diff --git a/tests/python/unittest/test_gluon.py b/tests/python/unittest/test_gluon.py index 952fdf7e366d..d91b3f02cd31 100644 --- a/tests/python/unittest/test_gluon.py +++ b/tests/python/unittest/test_gluon.py @@ -717,8 +717,8 @@ def test_lambda(): input_data = mx.nd.random.uniform(shape=(2, 3, 5, 7)) out1, out2, out3 = net1(input_data), net2(input_data), net3(input_data) - assert_almost_equal(out1.asnumpy(), out2.asnumpy()) - assert_almost_equal(out1.asnumpy(), out3.asnumpy()) + assert_almost_equal(out1.asnumpy(), out2.asnumpy(), rtol=1e-3) + assert_almost_equal(out1.asnumpy(), out3.asnumpy(), rtol=1e-3) @with_seed()