[MXNET-1234] Fix shape inference problems in Activation backward

[larroy]

Description

Exposes ReLU activation shape inference bug via C++ unit test.
Fixes Activation gradient shape inference problems due to different activation functions have different input parameters in backward, also depending on GPU / MKLDNN / CPU.
The logic now is now easier to maintain and safer.
This bug is uncovered when running C++ unit tests without MKL and without GPU, which we were not doing.

Fixes #13333

Checklist

Essentials

Please feel free to remove inapplicable items for your PR.

• The PR title starts with [MXNET-$JIRA_ID], where $JIRA_ID refers to the relevant JIRA issue created (except PRs with tiny changes)
• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage:
• Unit tests are added for small changes to verify correctness (e.g. adding a new operator)
• Nightly tests are added for complicated/long-running ones (e.g. changing distributed kvstore)
• Build tests will be added for build configuration changes (e.g. adding a new build option with NCCL)
• Code is well-documented:
• For user-facing API changes, API doc string has been updated.
• For new C++ functions in header files, their functionalities and arguments are documented.
• For new examples, README.md is added to explain the what the example does, the source of the dataset, expected performance on test set and reference to the original paper if applicable
• Check the API doc at http://mxnet-ci-doc.s3-accelerate.dualstack.amazonaws.com/PR-$PR_ID/$BUILD_ID/index.html
• To the my best knowledge, examples are either not affected by this change, or have been fixed to be compatible with this change

Checklist

Essentials

Please feel free to remove inapplicable items for your PR.

• The PR title starts with [MXNET-$JIRA_ID], where $JIRA_ID refers to the relevant JIRA issue created (except PRs with tiny changes)
• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage:
• Unit tests are added for small changes to verify correctness (e.g. adding a new operator)
• Nightly tests are added for complicated/long-running ones (e.g. changing distributed kvstore)
• Build tests will be added for build configuration changes (e.g. adding a new build option with NCCL)
• Code is well-documented:
• For user-facing API changes, API doc string has been updated.
• For new C++ functions in header files, their functionalities and arguments are documented.
• For new examples, README.md is added to explain the what the example does, the source of the dataset, expected performance on test set and reference to the original paper if applicable
• Check the API doc at http://mxnet-ci-doc.s3-accelerate.dualstack.amazonaws.com/PR-$PR_ID/$BUILD_ID/index.html
• To the my best knowledge, examples are either not affected by this change, or have been fixed to be compatible with this change

[stu1130]

@mxnet-label-bot add [pr-work-in-progress]
Thanks for doing this @larroy

[larroy]

@eric-haibin-lin @KellenSunderland Please review

[eric-haibin-lin]

Thanks for fixing this! Do we have an unit test that can be used to verify this bug?

[larroy]

@eric-haibin-lin It's in the patch as stated in the description and commit messages. The first commit introduces a failing test. If you have a better idea / suggestion let me know.

/~https://github.com/apache/incubator-mxnet/pull/13409/files#diff-91b46355311af5a0cdff26d6e706403fR43

[eric-haibin-lin]

Can we revert this change? The common practice is that if we change the elements in the vector, we will use pointer instead of const reference.

[larroy]

Right. We are not changing the vector in this function if you look closely. So your comment doesn't apply.

[eric-haibin-lin]

Oh right, I missed the cpp unit test. I thought it is going to be a python test. Nice work! One comment regarding coding style, otherwise looks good.

[lebeg]

Nice fix

[TaoLv]

If USE_CUDNN=1, then MKL-DNN code in this file will not be compiled?

[TaoLv]

If MXNet is built with CUDNN and MKL-DNN enabled at the same time, I hope the behavior is the activation will be dispatched to CUDNN if model context is gpu, otherwise it should be dispatched to MKL-DNN implementation. But with this change, MKL-DNN implementation will not be compiled.

[larroy]

This is a good point, the latest commit is just a test. What I observed is that we are calculating the backward in MKLDNN even if we have CUDNN.
via ActivationGradComputeExCPU as per:

/~https://github.com/apache/incubator-mxnet/blob/master/src/operator/nn/activation.cc#L189

Should't backward go through ActivationGradCompute and dispatched to cpu and GPU as you suggest? in this case we would need further changes to call mkldnn from ActivationGradCompute. Or is there something I'm missing?

[TaoLv]

I‘m not sure. But I think if MKL-DNN is enabled, the execution should be dispatched to FComputeEx according to the infer storage, but seems FComputeEx is not registered for gpu Activation. Need inputs from expert of execution engine.

[larroy]

Seems FComputeEx takes precedence:

/~https://github.com/apache/incubator-mxnet/blob/master/src/imperative/imperative.cc#L65

@eric-haibin-lin @piiswrong

[larroy]

My commit is passing all the unit tests, as it doesn't use MKLDNN if CUDNN is enabled. I think is the right way to go as of now, since we can't selectively dispatch to MKLDNN depending on the context without a bigger refactor.

[larroy]

@TaoLv can you approve? or have any other concerns / comments?

[TaoLv]

If it's a true problem for execution engine, I'm afraid same issue would also exist in other MKL-DNN operators. Also this workaround will hurt the performance when the users run model with CPU context and mxnet-cuxxmkl package.

@pengzhao-intel thoughts?

[pengzhao-intel]

@larroy Thanks for the fix :) The current solution can work and resolve the problem.

However, two things I want to highlight:

• Divergence
  this kind of workaround will lead divergent codes which are hard to understand and maintain later
• Performance degradation
  As @TaoLv mentioned, definitely, the performance will drop in CUXXMKL build compared with pure MKL build. And it's really hard to debug for end-user too.

So, I think we can work together to make

• An testcase/example to reproduce the issue.
  We still need to futher analysis and see if there're other solutions (maybe inside MKLDNN OP).
• A plan for the executor improvements
  Agree @larroy's comments, this would be a big refactor or chanllange so it's better to have a plan earily.

Feel free to correct me :)

[TaoLv]

@larroy Could you provide a simple reproducer for this problem? We can help to take a look. Actually, I cannot reproduce it with below code snippet:

import mxnet as mx
import numpy as np
np.random.seed(12345)

num_filter = 256
num_group = 1
kernel = (3, 3)
pad = (1, 1)
shape = (1, 256, 200, 233)

x = mx.sym.Variable('x')
w = mx.sym.Variable('w')

conv = mx.sym.Convolution(data=x, weight=w, num_filter=num_filter, num_group=num_group, kernel=kernel, no_bias=True, pad=pad)
relu = mx.sym.Activation(data=conv, act_type='relu', name='relu')
exe = relu.simple_bind(ctx=mx.gpu(), x=shape)

exe.arg_arrays[0][:] = np.random.normal(size=exe.arg_arrays[0].shape)
exe.arg_arrays[1][:] = np.random.normal(size=exe.arg_arrays[1].shape)

for i in range(10):
    exe.forward(is_train=True)
    exe.backward(exe.outputs[0])
    o = exe.grad_arrays[0]
    t = o.asnumpy()

[larroy]

@TaoLv the test that reproduces the problem is in the first commit in this PR. I propose to fix for now as this bug affects builds without mkl and ARM. And address the corner case performance problem in a subsequent PR.

[larroy]

Just to be clear. Compile CPU without mkl and run cppunit tests to exercise the bug. Also see my first commit

[larroy]

Do you guys have a better suggestion? I don't think with the current architecture there's a better solution, as there's no context information in functions like:

mxnet::op::ActivationGrad::operator() activation.cc:148 to know if we are in CPU or GPU context. As MKLDNN, CUDNN and CPU have all different set of inputs for Activation Backward, I don't see another possibility other than this fix. It's also a bit of a corner case having a CPU activation backward running when CUDNN is available, I don't see the concern when most users won't have such a case. When you have a GPU available why would you run it on the CPU?

[TaoLv]

I tried this test case on my machine with both MKL-DNN and CUDNN enabled, without the last commit. All the 5 activations run into MXNet CPU implementation (ActivationCompute<cpu> and ActivationGradCompute<cpu>). If set isGPU=true, all the 5 activations run into GPU implementation (ActivationCompute<gpu> and ActivationGradCompute<gpu>). Seems no MKL-DNN implementation was triggered for both isGPU=false and isGPU=true.

[TaoLv]

@larroy How to reproduce "What I observed is that we are calculating the backward in MKLDNN even if we have CUDNN." ?

[larroy]

@TaoLv you are right, this statement is not true in the master branch. I retract it.

We can use this patch to reproduce the bug and see what runs and doesn't run:

then run:

build/tests/mxnet_unit_tests --gtest_filter="ACTIVATION_PERF.ExecuteBidirectional"

diff --git a/src/operator/nn/activation-inl.h b/src/operator/nn/activation-inl.h
index 2705177..c656e33 100644
--- a/src/operator/nn/activation-inl.h
+++ b/src/operator/nn/activation-inl.h
@@ -88,6 +88,7 @@ void ActivationForward(const OpContext &ctx, const TBlob &in_data,
                        const OpReqType &req, const TBlob &out_data) {
   using namespace mshadow;
   using namespace mshadow::expr;
+  std::cout << "ActivationForward" << std::endl;
   Stream<xpu> *s = ctx.get_stream<xpu>();
   const size_t sz = in_data.shape_.Size();
   if (sz) {
@@ -104,6 +105,7 @@ template<typename xpu, typename ForwardOp, typename BackwardOp>
 void ActivationBackward(const OpContext &ctx, const TBlob &out_grad,
                         const TBlob &out_data, const OpReqType &req,
                         const TBlob &in_grad) {
+  std::cout << "ActivationBackward" << std::endl;
   using namespace mshadow;
   using namespace mshadow::expr;
   Stream<xpu> *s = ctx.get_stream<xpu>();
@@ -123,6 +125,7 @@ template<typename xpu>
 void ActivationComputeImpl(const nnvm::NodeAttrs& attrs, const OpContext &ctx,
                            const std::vector<TBlob>& inputs, const std::vector<OpReqType>& req,
                            const std::vector<TBlob>& outputs) {
+  std::cout << "ActivationComputeImpl" << std::endl;
   const ActivationParam& param = nnvm::get<ActivationParam>(attrs.parsed);
   switch (param.act_type) {
     case activation::kReLU:
@@ -154,6 +157,7 @@ template<typename xpu>
 void ActivationGradComputeImpl(const nnvm::NodeAttrs& attrs, const OpContext &ctx,
                            const std::vector<TBlob>& inputs, const std::vector<OpReqType>& req,
                            const std::vector<TBlob>& outputs) {
+  std::cout << "ActivationGradComputeImpl" << std::endl;
   const ActivationParam& param = nnvm::get<ActivationParam>(attrs.parsed);
   switch (param.act_type) {
     case activation::kReLU:
@@ -187,6 +191,7 @@ void ActivationCompute(const nnvm::NodeAttrs& attrs,
                        const std::vector<TBlob>& inputs,
                        const std::vector<OpReqType>& req,
                        const std::vector<TBlob>& outputs) {
+  std::cout << "ActivationCompute" << std::endl;
   CHECK_EQ(inputs.size(), 1U);
   CHECK_EQ(outputs.size(), 1U);
   ActivationComputeImpl<xpu>(attrs, ctx, inputs, req, outputs);
@@ -198,6 +203,7 @@ void ActivationGradCompute(const nnvm::NodeAttrs& attrs,
                            const std::vector<TBlob>& inputs,
                            const std::vector<OpReqType>& req,
                            const std::vector<TBlob>& outputs) {
+  std::cout << "ActivationGradCompute" << std::endl;
   const ActivationParam& param = nnvm::get<ActivationParam>(attrs.parsed);
 #if (MXNET_USE_CUDNN == 1 || MXNET_USE_MKLDNN == 1)
   bool relu = param.act_type == activation::kReLU;
diff --git a/src/operator/nn/activation.cc b/src/operator/nn/activation.cc
index ba44ebd..cb5d5ab 100644
--- a/src/operator/nn/activation.cc
+++ b/src/operator/nn/activation.cc
@@ -72,6 +72,7 @@ static void ActivationComputeExCPU(const nnvm::NodeAttrs& attrs,
                                    const std::vector<NDArray>& inputs,
                                    const std::vector<OpReqType>& req,
                                    const std::vector<NDArray>& outputs) {
+  std::cout << "ActivationComputeExCPU" << std::endl;
   CHECK_EQ(inputs.size(), 1U);
   CHECK_EQ(outputs.size(), 1U);
   if (SupportMKLDNN(inputs[0])) {
@@ -88,6 +89,7 @@ void ActivationGradComputeExCPU(const nnvm::NodeAttrs& attrs,
                                 const std::vector<NDArray>& inputs,
                                 const std::vector<OpReqType>& req,
                                 const std::vector<NDArray>& outputs) {
+  std::cout << "ActivationGradComputeExCPU" << std::endl;
   const ActivationParam& param = nnvm::get<ActivationParam>(attrs.parsed);
   bool relu = param.act_type == activation::kReLU;
   CHECK_EQ(inputs.size(), relu ? 2U : 3U);
diff --git a/src/operator/nn/activation.cu b/src/operator/nn/activation.cu
index 8892cc3..5aaeb78 100644
--- a/src/operator/nn/activation.cu
+++ b/src/operator/nn/activation.cu
@@ -51,6 +51,7 @@ void ActivationCompute<gpu>(const nnvm::NodeAttrs& attrs,
     const std::vector<TBlob>& inputs,
     const std::vector<OpReqType>& req,
     const std::vector<TBlob>& outputs) {
+  std::cout << "ActivationCompute (GPU)" << std::endl;
   CHECK_EQ(inputs.size(), 1U);
   CHECK_EQ(outputs.size(), 1U);
   const ActivationParam& param = nnvm::get<ActivationParam>(attrs.parsed);
@@ -75,6 +76,7 @@ void ActivationGradCompute<gpu>(const nnvm::NodeAttrs& attrs,
                                 const std::vector<TBlob>& inputs,
                                 const std::vector<OpReqType>& req,
                                 const std::vector<TBlob>& outputs) {
+  std::cout << "ActivationGradCompute (GPU)" << std::endl;
   const ActivationParam& param = nnvm::get<ActivationParam>(attrs.parsed);
   bool relu = param.act_type == activation::kReLU;
   CHECK_EQ(inputs.size(), relu ? 2U : 3U);
diff --git a/tests/cpp/operator/activation_perf.cc b/tests/cpp/operator/activation_perf.cc
index 1bd8ca8..bba8a3e 100644
--- a/tests/cpp/operator/activation_perf.cc
+++ b/tests/cpp/operator/activation_perf.cc
@@ -38,13 +38,27 @@ const kwargs_t basic_activation_args = { };
  * \brief Generic bidirectional sanity test
  */
 TEST(ACTIVATION_PERF, ExecuteBidirectional) {
+  using namespace std;
   TShape shape({5, 5});
-  kwargs_t kwargs = basic_activation_args;
-  kwargs.push_back({"act_type", "tanh"});
-
-  test::op::CoreOperatorRunner<float> runner;
-  runner.RunBidirectional(false, { shape }, test::op::CoreOpExecutor<float>::ArgsWithOpName(
-          kwargs, "Activation", "_backward_Activation"), 1);
+  vector<string> activations = {
+    "relu",
+    "sigmoid",
+    "tanh",
+    "softrelu",
+    "softsign"
+  };
+  for (const string& activation : activations) {
+    kwargs_t activation_args = {{"act_type", activation}};
+    test::op::CoreOperatorRunner<float> runner;
+    runner.RunBidirectional(false, { shape }, test::op::CoreOpExecutor<float>::ArgsWithOpName(
+            activation_args, "Activation", "_backward_Activation"), 1);
+  }
+  for (const string& activation : activations) {
+    kwargs_t activation_args = {{"act_type", activation}};
+    test::op::CoreOperatorRunner<float> runner;
+    runner.RunBidirectional(true, { shape }, test::op::CoreOpExecutor<float>::ArgsWithOpName(
+            activation_args, "Activation", "_backward_Activation"), 1);
+  }
 }
 
 /*!

[TaoLv]

Applied this patch to mxnet master branch and got below error:

[lvtao@mlt-gpu207 mxnet-official]$ ./mxnet_unit_tests --gtest_filter="ACTIVATION_PERF.ExecuteBidirectional"
Found CUDA Device #: 0 properties: 6.0

Note: Google Test filter = ACTIVATION_PERF.ExecuteBidirectional
[==========] Running 1 test from 1 test case.
[----------] Global test environment set-up.
[----------] 1 test from ACTIVATION_PERF
[ RUN      ] ACTIVATION_PERF.ExecuteBidirectional
unknown file: Failure
C++ exception with description "[10:31:44] src/operator/tensor/./../elemwise_op_common.h:176: Check failed: in_attrs->size() == static_cast<size_t>(n_in) (2 vs. 3)  in operator

Stack trace returned 10 entries:
[bt] (0) ./mxnet_unit_tests(dmlc::StackTrace()+0x4a) [0x42f64a]
[bt] (1) ./mxnet_unit_tests(dmlc::LogMessageFatal::~LogMessageFatal()+0x21) [0x42fca1]
[bt] (2) lib/libmxnet.so(bool mxnet::op::ElemwiseType<3l, 1l>(nnvm::NodeAttrs const&, std::vector<int, std::allocator<int> >*, std::vector<int, std::allocator<int> >*)+0xe2) [0x7f1f0691bd52]
[bt] (3) ./mxnet_unit_tests() [0x477593]
[bt] (4) ./mxnet_unit_tests() [0x43a007]
[bt] (5) ./mxnet_unit_tests() [0x479a20]
[bt] (6) ./mxnet_unit_tests() [0x5e4838]
[bt] (7) ./mxnet_unit_tests() [0x5e5bc8]
[bt] (8) ./mxnet_unit_tests() [0x5e6119]
[bt] (9) ./mxnet_unit_tests() [0x632904]

" thrown in the test body.
[  FAILED  ] ACTIVATION_PERF.ExecuteBidirectional (27 ms)
[----------] 1 test from ACTIVATION_PERF (27 ms total)

[----------] Global test environment tear-down
[==========] 1 test from 1 test case ran. (27 ms total)
[  PASSED  ] 0 tests.
[  FAILED  ] 1 test, listed below:
[  FAILED  ] ACTIVATION_PERF.ExecuteBidirectional

 1 FAILED TEST

[TaoLv]

I don't see the concern when most users won't have such a case. When you have a GPU available why would you run it on the CPU?

Not necessarily true. Otherwise, there is no need for MXNet to provide a mxnet-cuxxmkl pip package.

[larroy]

@TaoLv When I run your python reproducer in master with my patch printing which activation runs I get everything run on CPU:

ActivationForward
ActivationGradCompute
ActivationGradComputeImpl
ActivationBackward

Is this what you would expect?

I'm in commit 9979c3c with my previous patch applied to print activations.

[larroy]

This never executes on MKL:

Only if you bind the relu to cpu() context it does.

import mxnet as mx
import numpy as np
np.random.seed(12345)

num_filter = 256
num_group = 1
kernel = (3, 3)
pad = (1, 1)
shape = (1, 256, 200, 233)

x = mx.sym.Variable('x')
w = mx.sym.Variable('w')

conv = mx.sym.Convolution(data=x, weight=w, num_filter=num_filter, num_group=num_group, kernel=kernel, no_bias=True, pad=pad)
relu = mx.sym.Activation(data=conv, act_type='relu', name='relu')
exe = relu.simple_bind(ctx=mx.gpu(), x=shape)

exe.arg_arrays[0][:] = mx.nd.array(np.random.normal(size=exe.arg_arrays[0].shape), ctx=mx.cpu())
exe.arg_arrays[1][:] = mx.nd.array(np.random.normal(size=exe.arg_arrays[1].shape), ctx=mx.cpu())

for i in range(10):
    exe.forward(is_train=True)
    exe.backward(exe.outputs[0])
    o = exe.grad_arrays[0]
    t = o.asnumpy()

Output is:

ActivationGradCompute
ActivationGradComputeImpl
ActivationBackward
ActivationCompute
ActivationComputeImpl
ActivationForward
ActivationGradCompute
ActivationGradComputeImpl
ActivationBackward

[larroy]

@TaoLv @pengzhao-intel please check if the latest patch addresses your concerns.

[larroy]

One more clarification, shape inference doesn't fail in python because it doesn't call the operator using imperative. Such as Imperative::Invoke -> MXImperativeInvokeImpl which is going to call SetShapeType and thus fail on shape inference.

[TaoLv]

I'm validating benchmark performance with this PR. Will update soon.

[TaoLv]

I built MXNet from source with both CUDNN and MKL-DNN eanbled and tested the performance of AlexNet/VGG-16/ResNet-50/ResNet-152 from benchmark_score.py and didn't observe any regression. So I think my concern was addressed. Thank you for the contribution, @larroy .

[larroy]

@mxnet-label-bot [pr-awaiting-merge]