Merge ea45607 into f75a88c

mmcv/runner/base_runner.py

@@ -1,4 +1,5 @@
 # Copyright (c) Open-MMLab. All rights reserved.
+import copy
 import logging
 import os.path as osp
 import warnings
@@ -11,7 +12,7 @@
 from ..parallel import is_module_wrapper
 from .checkpoint import load_checkpoint
 from .dist_utils import get_dist_info
-from .hooks import HOOKS, Hook, IterTimerHook
+from .hooks import HOOKS, Hook
 from .log_buffer import LogBuffer
 from .priority import get_priority
 from .utils import get_time_str
@@ -414,12 +415,23 @@ def register_logger_hooks(self, log_config):
                 info, HOOKS, default_args=dict(interval=log_interval))
             self.register_hook(logger_hook, priority='VERY_LOW')
 
+    def register_timer_hook(self, timer_config):
+        if timer_config is None:
+            return
+        if isinstance(timer_config, dict):
+            timer_config_ = copy.deepcopy(timer_config)
+            hook = mmcv.buid_from_cfg(timer_config_, HOOKS)
+        else:
+            hook = timer_config
+        self.register_hook(hook)
+
     def register_training_hooks(self,
                                 lr_config,
                                 optimizer_config=None,
                                 checkpoint_config=None,
                                 log_config=None,
-                                momentum_config=None):
+                                momentum_config=None,
+                                timer_config=dict(type='IterTimerHook')):
         """Register default hooks for training.
 
         Default hooks include:
@@ -435,5 +447,5 @@ def register_training_hooks(self,
         self.register_momentum_hook(momentum_config)
         self.register_optimizer_hook(optimizer_config)
         self.register_checkpoint_hook(checkpoint_config)
-        self.register_hook(IterTimerHook())
+        self.register_timer_hook(timer_config)
         self.register_logger_hooks(log_config)