[incubator-mxnet] branch master updated: Fix a bug in sparse embedding operator (#11231)

[zh...@apache.org]

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 4da51b2  Fix a bug in sparse embedding operator (#11231)
4da51b2 is described below

commit 4da51b243e6f08bf53d27d96279998c8eaa5f039
Author: Haibin Lin <li...@gmail.com>
AuthorDate: Mon Jun 11 20:41:24 2018 -0700

    Fix a bug in sparse embedding operator (#11231)
    
    * add sparse block
    
    * add sparse embedding
    
    * add doc
    
    * lint
    
    * remove sparseblock
    
    * fix embedding
---
 src/operator/tensor/indexing_op.cu  |  3 ++-
 tests/python/unittest/test_gluon.py | 14 ++++++++++++++
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/src/operator/tensor/indexing_op.cu b/src/operator/tensor/indexing_op.cu
index 593593a..b0ee05e 100644
--- a/src/operator/tensor/indexing_op.cu
+++ b/src/operator/tensor/indexing_op.cu
@@ -188,7 +188,8 @@ void SparseEmbeddingDeterministicKernelLaunch(const OpContext& ctx,
   // estimate unique temp space
   IType* data_ptr = data.dptr<IType>();
   size_t *null_ptr = nullptr;
-  cub::DeviceSelect::Unique(NULL, unique_workspace_bytes, data_ptr, data_ptr,
+  // unique operations will be applied on sorted data
+  cub::DeviceSelect::Unique(NULL, unique_workspace_bytes, sorted_data, sorted_data,
     null_ptr, data_size, Stream<gpu>::GetStream(s));
   // One more space reserved for unique count
   size_t temp_workspace_bytes = std::max(unique_workspace_bytes,
diff --git a/tests/python/unittest/test_gluon.py b/tests/python/unittest/test_gluon.py
index bf1e0de..ced3063 100644
--- a/tests/python/unittest/test_gluon.py
+++ b/tests/python/unittest/test_gluon.py
@@ -753,8 +753,22 @@ def test_embedding():
             y.backward()
         assert (layer.weight.grad().asnumpy()[:5] == 1).all()
         assert (layer.weight.grad().asnumpy()[5:] == 0).all()
+
+    def check_embedding_large_input(sparse_grad):
+        embedding = mx.gluon.nn.Embedding(10, 1, sparse_grad=True)
+        embedding.initialize()
+        embedding.hybridize()
+        shape = (20481,)
+        with mx.autograd.record():
+            emb_in = embedding(mx.nd.ones(shape))
+            loss = emb_in.sum()
+        loss.backward()
+        assert embedding.weight.grad().data.sum().asscalar() == 20481
+
     check_embedding(True)
     check_embedding(False)
+    check_embedding_large_input(True)
+    check_embedding_large_input(False)
 
 @with_seed()
 def test_export():

-- 
To stop receiving notification emails like this one, please contact
zhasheng@apache.org.