You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@mxnet.apache.org by ma...@apache.org on 2020/12/10 06:35:35 UTC

[incubator-mxnet] branch master updated: [BUGFIX] Fix backward pass for nested CachedOp (#19614)

This is an automated email from the ASF dual-hosted git repository.

manuseth 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 63f2b32  [BUGFIX] Fix backward pass for nested CachedOp (#19614)
63f2b32 is described below

commit 63f2b3213f7b085d2198330f52c447e7e54d3383
Author: waytrue17 <52...@users.noreply.github.com>
AuthorDate: Wed Dec 9 22:33:36 2020 -0800

    [BUGFIX] Fix backward pass for nested CachedOp (#19614)
    
    * fix InvokeOperator
    
    * fix typo
    
    * update unsupported ops
    
    * remove contain_dynamic_shape
    
    Co-authored-by: Ubuntu <ub...@ip-172-31-60-22.us-west-2.compute.internal>
---
 python/mxnet/gluon/block.py                        |  6 ++---
 src/imperative/imperative_utils.cc                 |  2 +-
 .../subgraph/static_shape_subgraph_property.cc     |  6 ++---
 tests/python/unittest/test_dynamic_shape.py        | 26 +++++++++++++++++++++-
 4 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/python/mxnet/gluon/block.py b/python/mxnet/gluon/block.py
index 6b280e8..28e6f36 100644
--- a/python/mxnet/gluon/block.py
+++ b/python/mxnet/gluon/block.py
@@ -915,7 +915,7 @@ class HybridBlock(Block):
         self._monitor_all = False
         self._backend = None
         self._backend_opts = {}
-        self._partition_if_dynamic = False
+        self._partition_if_dynamic = True
         self._first_forward = True
 
     def __setattr__(self, name, value):
@@ -1174,7 +1174,7 @@ class HybridBlock(Block):
         return _regroup(out, self._out_format)
 
     def optimize_for(self, x, *args, backend=None, clear=False,
-                     partition_if_dynamic=False,
+                     partition_if_dynamic=True,
                      static_alloc=False,
                      static_shape=False,
                      inline_limit=2,
@@ -1281,7 +1281,7 @@ class HybridBlock(Block):
         self._clear_cached_op()
 
     def hybridize(self, active=True,
-                  partition_if_dynamic=False,
+                  partition_if_dynamic=True,
                   static_alloc=False,
                   static_shape=False,
                   inline_limit=2,
diff --git a/src/imperative/imperative_utils.cc b/src/imperative/imperative_utils.cc
index 654c0d3..9d15084 100644
--- a/src/imperative/imperative_utils.cc
+++ b/src/imperative/imperative_utils.cc
@@ -84,7 +84,7 @@ void InvokeOperator(const nnvm::IndexedGraph& idx,
   std::vector<uint32_t> &ref_count = *p_ref_count;
 
   const nnvm::IndexedGraph::Node& node = idx[node_idx];
-  if (node.source->op() == bwd_cached_op) {
+  if (node.source->op() == bwd_cached_op && node.source->attrs.name == "_cachedop_backward") {
     const auto& cached_op = dmlc::get<CachedOpPtr>(node.source->attrs.parsed);
     nnvm::Node* fwd_node = node.source->control_deps[0].get();
     auto fwd_node_id = idx.node_id(fwd_node);
diff --git a/src/operator/subgraph/static_shape_subgraph_property.cc b/src/operator/subgraph/static_shape_subgraph_property.cc
index 9dee7a6..923f709 100644
--- a/src/operator/subgraph/static_shape_subgraph_property.cc
+++ b/src/operator/subgraph/static_shape_subgraph_property.cc
@@ -54,10 +54,8 @@ class StaticShapeOpSelector: public SubgraphSelector {
   }
 
  private:
-    // static shape ops that fail backward pass inside subgraph CachedOp
-    // GitHub issue: https://github.com/apache/incubator-mxnet/issues/18823
-    std::unordered_set<std::string> unsupported_op_names_ {"Reshape", "_np_reshape", "transpose",
-                                                           "_npi_dstack", "_npi_hstack"};
+    // Rejecte the ops that have FInferShape registered but return true by CheckDynamicShapeExists()
+    std::unordered_set<std::string> unsupported_op_names_ {"_npi_dstack", "_npi_hstack"};
 };
 
 /*
diff --git a/tests/python/unittest/test_dynamic_shape.py b/tests/python/unittest/test_dynamic_shape.py
index e8bba4a..44d9010 100644
--- a/tests/python/unittest/test_dynamic_shape.py
+++ b/tests/python/unittest/test_dynamic_shape.py
@@ -121,4 +121,28 @@ def test_dynamic_shape_switch_hybridize():
 
     block.hybridize(static_alloc=True, static_shape=True)
     result = block(data, index)
-    assert_almost_equal(result.asnumpy(), result_nd)
\ No newline at end of file
+    assert_almost_equal(result.asnumpy(), result_nd)
+
+def test_dynamic_shape_backward():
+    # test dynamic shape ops with backward prop
+    class _TestBlock(gluon.HybridBlock):
+        def __init__(self):
+            super(_TestBlock, self).__init__()
+
+        def hybrid_forward(self, F, data, index):
+            return F.contrib.boolean_mask(F.sum(F.transpose(data)), index)
+
+    block = _TestBlock()
+    for static_alloc in [True, False]:
+      block.hybridize(static_alloc=static_alloc)
+      data = mx.nd.array([[1, 2, 3],[4, 5, 6],[7, 8, 9]])
+      index = mx.nd.array([1])
+      data.attach_grad()
+      with mx.autograd.record():
+        result = block(data, index)
+      result.backward()
+      result_nd = np.array([45.])
+      data_grad_nd = np.array([[1., 1., 1.], [1., 1., 1.], [1., 1., 1.]])
+      assert_almost_equal(result.asnumpy(), result_nd)
+      assert_almost_equal(data.grad.asnumpy(), data_grad_nd)
+