LRScheduler with update_on_kvstore=False

apache · Mar 9, 2019 · 5d59336 · 5d59336
1 parent ed83071
commit 5d59336
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 5 deletions.
diff --git a/python/mxnet/gluon/trainer.py b/python/mxnet/gluon/trainer.py
@@ -241,10 +241,10 @@ def _init_kvstore(self):
                 kvstore.set_optimizer(self._optimizer)
             self._kvstore = kvstore
             self._update_on_kvstore = update_on_kvstore
-            if self._optimizer.lr_scheduler and not self._update_on_kvstore:
-                raise ValueError("update_on_kvstore=False does not support " \
-                                 "optimizer with LRScheduler. Please " \
-                                 "consider setting learning rate manually.")
+            #if self._optimizer.lr_scheduler and not self._update_on_kvstore:
+            #    raise ValueError("update_on_kvstore=False does not support " \
+            #                     "optimizer with LRScheduler. Please " \
+            #                     "consider setting learning rate manually.")
         else:
             self._kvstore = None
             self._update_on_kvstore = None

diff --git a/python/mxnet/optimizer/optimizer.py b/python/mxnet/optimizer/optimizer.py
@@ -106,7 +106,8 @@ def __init__(self, rescale_grad=1., param_idx2name=None, wd=0.,
         self.wd_mult = {}
         self.begin_num_update = begin_num_update
         self.num_update = begin_num_update
-        self._index_update_count = {}
+        self._all_index_update_counts = {0 : {}}
+        self._index_update_count = self._all_index_update_counts[0]
         self.clip_gradient = clip_gradient
         self.multi_precision = multi_precision
         self.aggregate_num = 0
@@ -380,6 +381,11 @@ def set_wd_mult(self, args_wd_mult):
                     self.wd_mult[name] = float(attr[name]['__wd_mult__'])
         self.wd_mult.update(args_wd_mult)
 
+    def set_current_context(self, ctx):
+        if ctx not in self._all_index_update_counts:
+            self._all_index_update_counts[ctx] = {}
+        self._index_update_count = self._all_index_update_counts[ctx]
+
     def _update_count(self, index):
         """Updates num_update.
 
@@ -395,6 +401,7 @@ def _update_count(self, index):
                 self._index_update_count[idx] = self.begin_num_update
             self._index_update_count[idx] += 1
             self.num_update = max(self._index_update_count[idx], self.num_update)
+        print(self.num_update)
 
     def _get_lrs(self, indices):
         """Gets the learning rates given the indices of the weights.
@@ -1623,6 +1630,8 @@ def __call__(self, index, grad, weight):
             indices = index
             grads = grad
             weights = weight
+        if weights:
+            self.optimizer.set_current_context(weights[0].context.device_id)
         for i, idx in enumerate(indices):
             # convert ctypes.char_p.value back to python str if needed
             if isinstance(idx, bytes):