[GitHub] piiswrong commented on a change in pull request #10861: split trainer.step into allreduce and update

[GitBox <gi...@apache.org>]

piiswrong commented on a change in pull request #10861: split trainer.step into allreduce and update
URL: https://github.com/apache/incubator-mxnet/pull/10861#discussion_r187444700
 
 

 ##########
 File path: python/mxnet/gluon/trainer.py
 ##########
 @@ -153,11 +161,55 @@ def set_learning_rate(self, lr):
         else:
             self._optimizer.set_learning_rate(lr)
 
-
     def step(self, batch_size, ignore_stale_grad=False):
         """Makes one step of parameter update. Should be called after
         `autograd.compute_gradient` and outside of `record()` scope.
 
+        Parameters
+        ----------
+        batch_size : int
+            Batch size of data processed. Gradient will be normalized by `1/batch_size`.
+            Set this to 1 if you normalized loss manually with `loss = mean(loss)`.
+        ignore_stale_grad : bool, optional, default=False
+            If true, ignores Parameters with stale gradient (gradient that has not
+            been updated by `backward` after last step) and skip update.
+        """
+        if not self._kv_initialized:
+            self._init_kvstore()
+
+        self._allreduce(batch_size, skip_pull=self._update_on_kvstore)
+        self.update(batch_size, ignore_stale_grad)
+
+    def allreduce(self, batch_size):
+        """For each parameter, reduce the gradients from different contexts. Should be called after
+        `autograd.compute_gradient`, outside of `record()` scope, and before `trainer.update`.
+        """
+        self._allreduce(batch_size, skip_pull=False)
+
+    def _allreduce(self, batch_size, skip_pull):
+        if not self._kv_initialized:
+            self._init_kvstore()
+
+        self._optimizer.rescale_grad = self._scale / batch_size
+
+        for i, param in enumerate(self._params):
+            if param.grad_req == 'null':
+                continue
+
+            if self._kvstore:
 
 Review comment:
   do this outside of for loop

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
users@infra.apache.org


With regards,
Apache Git Services