[incubator-mxnet] branch master updated: Fix gluon.Trainer regression if no kvstore is used with sparse gradients (#17199)

[la...@apache.org]

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

lausen 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 8dee5b7  Fix gluon.Trainer regression if no kvstore is used with sparse gradients (#17199)
8dee5b7 is described below

commit 8dee5b7bf0307450a8663b3ae4bad53564a08402
Author: Leonard Lausen <la...@amazon.com>
AuthorDate: Wed Jan 1 13:21:35 2020 +0000

    Fix gluon.Trainer regression if no kvstore is used with sparse gradients (#17199)
---
 python/mxnet/gluon/trainer.py               |  2 +-
 tests/python/unittest/test_gluon_trainer.py | 15 +++++++++++++++
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/python/mxnet/gluon/trainer.py b/python/mxnet/gluon/trainer.py
index 966ed2c..a27c951 100644
--- a/python/mxnet/gluon/trainer.py
+++ b/python/mxnet/gluon/trainer.py
@@ -221,7 +221,7 @@ class Trainer(object):
                                      "when sparse gradients are present.")
                 update_on_kvstore = config['update_on_kvstore']
             # raise err if a custom kvstore is used for sparse training
-            if not isinstance(kvstore, KVStore):
+            if kvstore is not None and not isinstance(kvstore, KVStore):
                 raise ValueError("Cannot use {} for multi-device training with sparse gradients"
                                  .format(type(kvstore)))
 
diff --git a/tests/python/unittest/test_gluon_trainer.py b/tests/python/unittest/test_gluon_trainer.py
index fbd04ee..350700c 100644
--- a/tests/python/unittest/test_gluon_trainer.py
+++ b/tests/python/unittest/test_gluon_trainer.py
@@ -47,6 +47,21 @@ def test_multi_trainer():
     trainer1 = gluon.Trainer([x], 'sgd')
 
 @with_seed()
+def test_trainer_with_sparse_grad_on_single_context():
+    x = gluon.Parameter('x', shape=(10,), grad_stype='row_sparse')
+    x.initialize(ctx=[mx.cpu(0)], init='zeros')
+    trainer = gluon.Trainer([x], 'sgd', {'learning_rate': 1.0, 'momentum': 0.5})
+    with mx.autograd.record():
+        for w in x.list_data():
+            y = w + 1
+            y.backward()
+    trainer.step(1)
+
+    assert trainer._update_on_kvstore is None
+    assert trainer._kvstore is None  # No kvstore created for single-device training
+    assert (x.data(mx.cpu(0)).asnumpy() == -1).all()
+
+@with_seed()
 def test_trainer_with_teststore():
     x = gluon.Parameter('x', shape=(10,))
     x.initialize(ctx=[mx.cpu(0), mx.cpu(1)], init='zeros')