add FP16 optimizer (#432)

mmcv/runner/__init__.py

@@ -4,6 +4,7 @@ from .checkpoint import (_load_checkpoint, load_checkpoint, load_state_dict,
                          save_checkpoint, weights_to_cpu)
 from .dist_utils import get_dist_info, init_dist, master_only
 from .epoch_based_runner import EpochBasedRunner, Runner
+from .fp16_utils import auto_fp16, force_fp32
 from .hooks import (HOOKS, CheckpointHook, ClosureHook, DistSamplerSeedHook,
                     Hook, IterTimerHook, LoggerHook, LrUpdaterHook,
                     MlflowLoggerHook, OptimizerHook, PaviLoggerHook,
@@ -27,5 +28,5 @@ __all__ = [
     'obj_from_dict', 'init_dist', 'get_dist_info', 'master_only',
     'OPTIMIZER_BUILDERS', 'OPTIMIZERS', 'DefaultOptimizerConstructor',
     'build_optimizer', 'build_optimizer_constructor', 'IterLoader',
-    'set_random_seed'
+    'set_random_seed', 'auto_fp16', 'force_fp32'
 ]

mmcv/runner/hooks/__init__.py

@@ -8,12 +8,13 @@ from .logger import (LoggerHook, MlflowLoggerHook, PaviLoggerHook,
 from .lr_updater import LrUpdaterHook
 from .memory import EmptyCacheHook
 from .momentum_updater import MomentumUpdaterHook
-from .optimizer import OptimizerHook
+from .optimizer import Fp16OptimizerHook, OptimizerHook
 from .sampler_seed import DistSamplerSeedHook
 
 __all__ = [
     'HOOKS', 'Hook', 'CheckpointHook', 'ClosureHook', 'LrUpdaterHook',
-    'OptimizerHook', 'IterTimerHook', 'DistSamplerSeedHook', 'EmptyCacheHook',
-    'LoggerHook', 'MlflowLoggerHook', 'PaviLoggerHook', 'TextLoggerHook',
-    'TensorboardLoggerHook', 'WandbLoggerHook', 'MomentumUpdaterHook'
+    'OptimizerHook', 'Fp16OptimizerHook', 'IterTimerHook',
+    'DistSamplerSeedHook', 'EmptyCacheHook', 'LoggerHook', 'MlflowLoggerHook',
+    'PaviLoggerHook', 'TextLoggerHook', 'TensorboardLoggerHook',
+    'WandbLoggerHook', 'MomentumUpdaterHook'
 ]

mmcv/runner/hooks/optimizer.py

@@ -31,6 +31,7 @@ class OptimizerHook(Hook):
         runner.optimizer.step()
 
 
+@HOOKS.register_module()
 class Fp16OptimizerHook(OptimizerHook):
     """FP16 optimizer hook.