add lr and total iter in trainer (tusen-ai#263)

Author: RogerChern

core/detection_module.py

@@ -23,6 +23,7 @@
 import time
 import logging
 import warnings
+from collections import namedtuple
 
 import mxnet
 
@@ -35,7 +36,6 @@
 from mxnet.module.executor_group import DataParallelExecutorGroup
 from mxnet.model import _create_kvstore, _initialize_kvstore, _update_params, _update_params_on_kvstore
 from mxnet.model import load_checkpoint
-from mxnet.model import BatchEndParam
 from mxnet.initializer import Uniform, InitDesc
 from mxnet.io import DataDesc
 from mxnet.ndarray import zeros
@@ -44,6 +44,10 @@
 from mxnet.module.module import Module
 
 
+BatchEndParam = namedtuple('BatchEndParams',
+                           ['epoch', 'nbatch', 'eval_metric', 'lr', 'iter', 'locals'])
+
+
 class DetModule(BaseModule):
     """Module is a basic module that wrap a `Symbol`. It is functionally the same
     as the `FeedForward` model, except under the module API.
@@ -1025,6 +1029,8 @@ def fit(self, train_data, eval_data=None, eval_metric='acc',
                 if batch_end_callback is not None:
                     batch_end_params = BatchEndParam(epoch=epoch, nbatch=nbatch,
                                                      eval_metric=eval_metric,
+                                                     lr=self._optimizer.lr_scheduler(total_iter),
+                                                     iter=total_iter,
                                                      locals=locals())
                     for callback in _as_list(batch_end_callback):
                         callback(batch_end_params)

detection_train.py

@@ -235,6 +235,7 @@ def train_net(config):
                 mode='cosine',
                 base_lr=base_lr,
                 target_lr=0,
+                offset=pOpt.warmup.iter,
                 niters=(iter_per_epoch * (end_epoch - begin_epoch)) - pOpt.warmup.iter
             )
             lr_scheduler = LRSequential([warmup_lr_scheduler, cosine_lr_scheduler])
@@ -248,6 +249,7 @@ def train_net(config):
                 mode='cosine',
                 base_lr=base_lr,
                 target_lr=0,
+                offset=pOpt.warmup.iter,
                 niters=iter_per_epoch * (end_epoch - begin_epoch)
             )
         else:

utils/callback.py

@@ -23,12 +23,13 @@ def __call__(self, param):
                 speed = self.frequent * self.batch_size / (time.time() - self.tic)
                 if param.eval_metric is not None:
                     name, value = param.eval_metric.get()
-                    s = "Epoch[%d] Batch [%d]\tSpeed: %.2f samples/sec\tTrain-" % (param.epoch, count, speed)
+                    s = "Epoch[%d] Batch [%d]\tIter: %d\tLr: %.5f\tSpeed: %.2f samples/sec\tTrain-" % \
+                        (param.epoch, count, param.iter, param.lr, speed)
                     for n, v in zip(name, value):
                         s += "%s=%f,\t" % (n, v)
                     logging.info(s)
                 else:
-                    logging.info("Iter[%d] Batch [%d]\tSpeed: %.2f samples/sec",
+                    logging.info("Epoch[%d] Batch [%d]\tSpeed: %.2f samples/sec",
                                  param.epoch, count, speed)
                 self.tic = time.time()
         else: