Skip to content

GitLab

Unverified Commit c77e95a6 authored by Ziyi Wu's avatar Ziyi Wu Committed by GitHub
Browse files

[Enhance] Add step momentum updater to support MMDet3D (#1011) 

* add StepMomentumUpdaterHook

* add unit test

* fix typos

* refactor step updater

* replace stage with exp

* fix linting error

* use all() operation
parent 1a5bf762
Hide whitespace changes
Inline Side-by-side
Showing with 116 additions and 49 deletions
mmcv/runner/hooks/lr_updater.py
View file @ c77e95a6
......@@ -2,6 +2,7 @@
import numbers
from math import cos, pi
import mmcv
from .hook import HOOKS, Hook
......@@ -177,10 +178,9 @@ class StepLrUpdaterHook(LrUpdaterHook):
"""
def __init__(self, step, gamma=0.1, min_lr=None, **kwargs):
assert isinstance(step, (list, int))
if isinstance(step, list):
for s in step:
assert isinstance(s, int) and s > 0
assert mmcv.is_list_of(step, int)
assert all([s > 0 for s in step])
elif isinstance(step, int):
assert step > 0
else:
......@@ -193,19 +193,17 @@ class StepLrUpdaterHook(LrUpdaterHook):
def get_lr(self, runner, base_lr):
progress = runner.epoch if self.by_epoch else runner.iter
# calculate exponential term
if isinstance(self.step, int):
lr = base_lr * (self.gamma**(progress // self.step))
if self.min_lr is not None:
# clip to a minimum value
lr = max(lr, self.min_lr)
return lr
exp = len(self.step)
for i, s in enumerate(self.step):
if progress < s:
exp = i
break
lr = base_lr * self.gamma**exp
exp = progress // self.step
else:
exp = len(self.step)
for i, s in enumerate(self.step):
if progress < s:
exp = i
break
lr = base_lr * (self.gamma**exp)
if self.min_lr is not None:
# clip to a minimum value
lr = max(lr, self.min_lr)
......
mmcv/runner/hooks/momentum_updater.py
View file @ c77e95a6
import mmcv
from .hook import HOOKS, Hook
from .lr_updater import annealing_cos, annealing_linear, format_param
......@@ -148,6 +149,54 @@ class MomentumUpdaterHook(Hook):
self._set_momentum(runner, warmup_momentum)
@HOOKS.register_module()
class StepMomentumUpdaterHook(MomentumUpdaterHook):
"""Step momentum scheduler with min value clipping.
Args:
step (int | list[int]): Step to decay the momentum. If an int value is
given, regard it as the decay interval. If a list is given, decay
momentum at these steps.
gamma (float, optional): Decay momentum ratio. Default: 0.5.
min_momentum (float, optional): Minimum momentum value to keep. If
momentum after decay is lower than this value, it will be clipped
accordingly. If None is given, we don't perform lr clipping.
Default: None.
"""
def __init__(self, step, gamma=0.5, min_momentum=None, **kwargs):
if isinstance(step, list):
assert mmcv.is_list_of(step, int)
assert all([s > 0 for s in step])
elif isinstance(step, int):
assert step > 0
else:
raise TypeError('"step" must be a list or integer')
self.step = step
self.gamma = gamma
self.min_momentum = min_momentum
super(StepMomentumUpdaterHook, self).__init__(**kwargs)
def get_momentum(self, runner, base_momentum):
progress = runner.epoch if self.by_epoch else runner.iter
# calculate exponential term
if isinstance(self.step, int):
exp = progress // self.step
else:
exp = len(self.step)
for i, s in enumerate(self.step):
if progress < s:
exp = i
break
momentum = base_momentum * (self.gamma**exp)
if self.min_momentum is not None:
# clip to a minimum value
momentum = max(momentum, self.min_momentum)
return momentum
@HOOKS.register_module()
class CosineAnnealingMomentumUpdaterHook(MomentumUpdaterHook):
......@@ -419,11 +468,12 @@ class OneCycleMomentumUpdaterHook(MomentumUpdaterHook):
end_iter = phase['end_iter']
if curr_iter <= end_iter or i == len(self.momentum_phases) - 1:
pct = (curr_iter - start_iter) / (end_iter - start_iter)
lr = self.anneal_func(param_group[phase['start_momentum']],
param_group[phase['end_momentum']], pct)
momentum = self.anneal_func(
param_group[phase['start_momentum']],
param_group[phase['end_momentum']], pct)
break
start_iter = end_iter
return lr
return momentum
def get_regular_momentum(self, runner):
if isinstance(runner.optimizer, dict):
......
tests/test_runner/test_hooks.py
View file @ c77e95a6
......@@ -540,7 +540,7 @@ def test_cosine_restart_lr_update_hook(multi_optimziers):
@pytest.mark.parametrize('multi_optimziers', (True, False))
def test_step_lr_update_hook(multi_optimziers):
def test_step_runner_hook(multi_optimziers):
"""Test StepLrUpdaterHook."""
with pytest.raises(TypeError):
# `step` should be specified
......@@ -557,6 +557,15 @@ def test_step_lr_update_hook(multi_optimziers):
loader = DataLoader(torch.ones((30, 2)))
runner = _build_demo_runner(multi_optimziers=multi_optimziers)
# add momentum scheduler
hook_cfg = dict(
type='StepMomentumUpdaterHook',
by_epoch=False,
step=5,
gamma=0.5,
min_momentum=0.05)
runner.register_hook_from_cfg(hook_cfg)
# add step LR scheduler
hook = StepLrUpdaterHook(by_epoch=False, step=5, gamma=0.5, min_lr=1e-3)
runner.register_hook(hook)
......@@ -583,36 +592,36 @@ def test_step_lr_update_hook(multi_optimziers):
'train', {
'learning_rate/model1': 0.01,
'learning_rate/model2': 0.005,
'momentum/model1': 0.95,
'momentum/model2': 0.9
'momentum/model1': 0.475,
'momentum/model2': 0.45
}, 6),
call(
'train', {
'learning_rate/model1': 0.0025,
'learning_rate/model2': 0.00125,
'momentum/model1': 0.95,
'momentum/model2': 0.9
'momentum/model1': 0.11875,
'momentum/model2': 0.1125
}, 16),
call(
'train', {
'learning_rate/model1': 0.00125,
'learning_rate/model2': 0.001,
'momentum/model1': 0.95,
'momentum/model2': 0.9
'momentum/model1': 0.059375,
'momentum/model2': 0.05625
}, 21),
call(
'train', {
'learning_rate/model1': 0.001,
'learning_rate/model2': 0.001,
'momentum/model1': 0.95,
'momentum/model2': 0.9
'momentum/model1': 0.05,
'momentum/model2': 0.05
}, 26),
call(
'train', {
'learning_rate/model1': 0.001,
'learning_rate/model2': 0.001,
'momentum/model1': 0.95,
'momentum/model2': 0.9
'momentum/model1': 0.05,
'momentum/model2': 0.05
}, 30)
]
else:
......@@ -623,23 +632,23 @@ def test_step_lr_update_hook(multi_optimziers):
}, 1),
call('train', {
'learning_rate': 0.01,
'momentum': 0.95
'momentum': 0.475
}, 6),
call('train', {
'learning_rate': 0.0025,
'momentum': 0.95
'momentum': 0.11875
}, 16),
call('train', {
'learning_rate': 0.00125,
'momentum': 0.95
'momentum': 0.059375
}, 21),
call('train', {
'learning_rate': 0.001,
'momentum': 0.95
'momentum': 0.05
}, 26),
call('train', {
'learning_rate': 0.001,
'momentum': 0.95
'momentum': 0.05
}, 30)
]
hook.writer.add_scalars.assert_has_calls(calls, any_order=True)
......@@ -649,6 +658,14 @@ def test_step_lr_update_hook(multi_optimziers):
loader = DataLoader(torch.ones((10, 2)))
runner = _build_demo_runner(multi_optimziers=multi_optimziers)
# add momentum scheduler
hook_cfg = dict(
type='StepMomentumUpdaterHook',
by_epoch=False,
step=[4, 6, 8],
gamma=0.1)
runner.register_hook_from_cfg(hook_cfg)
# add step LR scheduler
hook = StepLrUpdaterHook(by_epoch=False, step=[4, 6, 8], gamma=0.1)
runner.register_hook(hook)
......@@ -675,22 +692,22 @@ def test_step_lr_update_hook(multi_optimziers):
'train', {
'learning_rate/model1': 0.002,
'learning_rate/model2': 0.001,
'momentum/model1': 0.95,
'momentum/model2': 0.9
'momentum/model1': 9.5e-2,
'momentum/model2': 9.000000000000001e-2
}, 5),
call(
'train', {
'learning_rate/model1': 2.0000000000000004e-4,
'learning_rate/model2': 1.0000000000000002e-4,
'momentum/model1': 0.95,
'momentum/model2': 0.9
'momentum/model1': 9.500000000000001e-3,
'momentum/model2': 9.000000000000003e-3
}, 7),
call(
'train', {
'learning_rate/model1': 2.0000000000000005e-05,
'learning_rate/model2': 1.0000000000000003e-05,
'momentum/model1': 0.95,
'momentum/model2': 0.9
'momentum/model1': 9.500000000000002e-4,
'momentum/model2': 9.000000000000002e-4
}, 9)
]
else:
......@@ -701,16 +718,18 @@ def test_step_lr_update_hook(multi_optimziers):
}, 1),
call('train', {
'learning_rate': 0.002,
'momentum': 0.95
'momentum': 0.095
}, 5),
call('train', {
'learning_rate': 2.0000000000000004e-4,
'momentum': 0.95
}, 7),
call('train', {
'learning_rate': 2.0000000000000005e-05,
'momentum': 0.95
}, 9)
call(
'train', {
'learning_rate': 2.0000000000000004e-4,
'momentum': 9.500000000000001e-3
}, 7),
call(
'train', {
'learning_rate': 2.0000000000000005e-05,
'momentum': 9.500000000000002e-4
}, 9)
]
hook.writer.add_scalars.assert_has_calls(calls, any_order=True)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment