[Enhance] Refactor logger (#659)

* [Enhance] Refactor logger

* fixed test

* make commit optional

* remove debug info

* fixed test

mmcv/runner/hooks/logger/base.py

@@ -89,13 +89,6 @@ class LoggerHook(Hook):
             current_iter = runner.iter + 1
         return current_iter
 
-    def get_step(self, runner):
-        """Get the total training step/epoch."""
-        if self.get_mode(runner) == 'val' and self.by_epoch:
-            return self.get_epoch(runner)
-        else:
-            return self.get_iter(runner)
-
     def get_lr_tags(self, runner):
         tags = {}
         lrs = runner.current_lr()

mmcv/runner/hooks/logger/mlflow.py

@@ -69,7 +69,7 @@ class MlflowLoggerHook(LoggerHook):
     def log(self, runner):
         tags = self.get_loggable_tags(runner)
         if tags:
-            self.mlflow.log_metrics(tags, step=self.get_step(runner))
+            self.mlflow.log_metrics(tags, step=self.get_iter(runner))
 
     @master_only
     def after_run(self, runner):

mmcv/runner/hooks/logger/pavi.py

@@ -69,6 +69,13 @@ class PaviLoggerHook(LoggerHook):
         if self.add_graph:
             self.writer.add_graph(runner.model)
 
+    def get_step(self, runner):
+        """Get the total training step/epoch."""
+        if self.get_mode(runner) == 'val' and self.by_epoch:
+            return self.get_epoch(runner)
+        else:
+            return self.get_iter(runner)
+
     @master_only
     def log(self, runner):
         tags = self.get_loggable_tags(runner, add_mode=False)

mmcv/runner/hooks/logger/tensorboard.py

@@ -46,9 +46,9 @@ class TensorboardLoggerHook(LoggerHook):
         tags = self.get_loggable_tags(runner, allow_text=True)
         for tag, val in tags.items():
             if isinstance(val, str):
-                self.writer.add_text(tag, val, self.get_step(runner))
+                self.writer.add_text(tag, val, self.get_iter(runner))
             else:
-                self.writer.add_scalar(tag, val, self.get_step(runner))
+                self.writer.add_scalar(tag, val, self.get_iter(runner))
 
     @master_only
     def after_run(self, runner):

mmcv/runner/hooks/logger/wandb.py

@@ -12,11 +12,13 @@ class WandbLoggerHook(LoggerHook):
                  interval=10,
                  ignore_last=True,
                  reset_flag=True,
+                 commit=True,
                  by_epoch=True):
         super(WandbLoggerHook, self).__init__(interval, ignore_last,
                                               reset_flag, by_epoch)
         self.import_wandb()
         self.init_kwargs = init_kwargs
+        self.commit = commit
 
     def import_wandb(self):
         try:
@@ -39,7 +41,8 @@ class WandbLoggerHook(LoggerHook):
     def log(self, runner):
         tags = self.get_loggable_tags(runner)
         if tags:
-            self.wandb.log(tags, step=self.get_step(runner))
+            self.wandb.log(
+                tags, step=self.get_iter(runner), commit=self.commit)
 
     @master_only
     def after_run(self, runner):

tests/test_runner/test_hooks.py

@@ -346,7 +346,7 @@ def test_mlflow_hook(log_model):
         {
             'learning_rate': 0.02,
             'momentum': 0.95
-        }, step=1)
+        }, step=6)
     if log_model:
         hook.mlflow_pytorch.log_model.assert_called_with(
             runner.model, 'models')
@@ -369,7 +369,8 @@ def test_wandb_hook():
         'learning_rate': 0.02,
         'momentum': 0.95
     },
-                                      step=1)
+                                      step=6,
+                                      commit=True)
     hook.wandb.join.assert_called_with()