Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Feature] Support for linearly learning rate decay #1627

Merged
merged 5 commits into from
Mar 31, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion docs/en/community/contributing.md
9 changes: 6 additions & 3 deletions mmcv/runner/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,10 +23,12 @@
CosineRestartLrUpdaterHook, CyclicLrUpdaterHook,
ExpLrUpdaterHook, FixedLrUpdaterHook,
FlatCosineAnnealingLrUpdaterHook,
InvLrUpdaterHook, LrUpdaterHook,
OneCycleLrUpdaterHook, PolyLrUpdaterHook)
InvLrUpdaterHook, LinearAnnealingLrUpdaterHook,
LrUpdaterHook, OneCycleLrUpdaterHook,
PolyLrUpdaterHook)
from .hooks.momentum_updater import (CosineAnnealingMomentumUpdaterHook,
CyclicMomentumUpdaterHook,
LinearAnnealingMomentumUpdaterHook,
MomentumUpdaterHook,
OneCycleMomentumUpdaterHook,
StepMomentumUpdaterHook)
Expand Down Expand Up @@ -62,5 +64,6 @@
'_load_checkpoint_with_prefix', 'EvalHook', 'DistEvalHook', 'Sequential',
'ModuleDict', 'ModuleList', 'GradientCumulativeOptimizerHook',
'GradientCumulativeFp16OptimizerHook', 'DefaultRunnerConstructor',
'SegmindLoggerHook'
'SegmindLoggerHook', 'LinearAnnealingMomentumUpdaterHook',
'LinearAnnealingLrUpdaterHook'
]
12 changes: 8 additions & 4 deletions mmcv/runner/hooks/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,11 +12,14 @@
CosineRestartLrUpdaterHook, CyclicLrUpdaterHook,
ExpLrUpdaterHook, FixedLrUpdaterHook,
FlatCosineAnnealingLrUpdaterHook, InvLrUpdaterHook,
LrUpdaterHook, OneCycleLrUpdaterHook,
PolyLrUpdaterHook, StepLrUpdaterHook)
LinearAnnealingLrUpdaterHook, LrUpdaterHook,
OneCycleLrUpdaterHook, PolyLrUpdaterHook,
StepLrUpdaterHook)
from .memory import EmptyCacheHook
from .momentum_updater import (CosineAnnealingMomentumUpdaterHook,
CyclicMomentumUpdaterHook, MomentumUpdaterHook,
CyclicMomentumUpdaterHook,
LinearAnnealingMomentumUpdaterHook,
MomentumUpdaterHook,
OneCycleMomentumUpdaterHook,
StepMomentumUpdaterHook)
from .optimizer import (Fp16OptimizerHook, GradientCumulativeFp16OptimizerHook,
Expand All @@ -39,5 +42,6 @@
'CyclicMomentumUpdaterHook', 'OneCycleMomentumUpdaterHook',
'SyncBuffersHook', 'EMAHook', 'EvalHook', 'DistEvalHook', 'ProfilerHook',
'GradientCumulativeOptimizerHook', 'GradientCumulativeFp16OptimizerHook',
'SegmindLoggerHook'
'SegmindLoggerHook', 'LinearAnnealingLrUpdaterHook',
'LinearAnnealingMomentumUpdaterHook'
]
40 changes: 40 additions & 0 deletions mmcv/runner/hooks/lr_updater.py
Original file line number Diff line number Diff line change
Expand Up @@ -256,6 +256,14 @@ def get_lr(self, runner, base_lr):

@HOOKS.register_module()
class CosineAnnealingLrUpdaterHook(LrUpdaterHook):
"""CosineAnnealing LR scheduler.

Args:
min_lr (float, optional): The minimum lr. Default: None.
min_lr_ratio (float, optional): The ratio of minimum lr to the base lr.
Either `min_lr` or `min_lr_ratio` should be specified.
Default: None.
"""

def __init__(self, min_lr=None, min_lr_ratio=None, **kwargs):
assert (min_lr is None) ^ (min_lr_ratio is None)
Expand Down Expand Up @@ -644,6 +652,38 @@ def get_lr(self, runner, base_lr):
return lr


@HOOKS.register_module()
class LinearAnnealingLrUpdaterHook(LrUpdaterHook):
zhouzaida marked this conversation as resolved.
Show resolved Hide resolved
"""Linear annealing LR Scheduler decays the learning rate of each parameter
group linearly.

Args:
min_lr (float, optional): The minimum lr. Default: None.
min_lr_ratio (float, optional): The ratio of minimum lr to the base lr.
Either `min_lr` or `min_lr_ratio` should be specified.
Default: None.
"""

def __init__(self, min_lr=None, min_lr_ratio=None, **kwargs):
assert (min_lr is None) ^ (min_lr_ratio is None)
self.min_lr = min_lr
self.min_lr_ratio = min_lr_ratio
super(LinearAnnealingLrUpdaterHook, self).__init__(**kwargs)

def get_lr(self, runner, base_lr):
if self.by_epoch:
progress = runner.epoch
max_progress = runner.max_epochs
else:
progress = runner.iter
max_progress = runner.max_iters
if self.min_lr_ratio is not None:
target_lr = base_lr * self.min_lr_ratio
else:
target_lr = self.min_lr
return annealing_linear(base_lr, target_lr, progress / max_progress)


def annealing_cos(start, end, factor, weight=1):
"""Calculate annealing cos learning rate.

Expand Down
42 changes: 42 additions & 0 deletions mmcv/runner/hooks/momentum_updater.py
Original file line number Diff line number Diff line change
Expand Up @@ -200,6 +200,15 @@ def get_momentum(self, runner, base_momentum):

@HOOKS.register_module()
class CosineAnnealingMomentumUpdaterHook(MomentumUpdaterHook):
"""Cosine annealing LR Momentum decays the Momentum of each parameter group
linearly.

Args:
min_momentum (float, optional): The minimum momentum. Default: None.
min_momentum_ratio (float, optional): The ratio of minimum momentum to
the base momentum. Either `min_momentum` or `min_momentum_ratio`
should be specified. Default: None.
"""

def __init__(self, min_momentum=None, min_momentum_ratio=None, **kwargs):
assert (min_momentum is None) ^ (min_momentum_ratio is None)
Expand All @@ -222,6 +231,39 @@ def get_momentum(self, runner, base_momentum):
progress / max_progress)


@HOOKS.register_module()
class LinearAnnealingMomentumUpdaterHook(MomentumUpdaterHook):
"""Linear annealing LR Momentum decays the Momentum of each parameter group
linearly.

Args:
min_momentum (float, optional): The minimum momentum. Default: None.
min_momentum_ratio (float, optional): The ratio of minimum momentum to
the base momentum. Either `min_momentum` or `min_momentum_ratio`
should be specified. Default: None.
"""

zhouzaida marked this conversation as resolved.
Show resolved Hide resolved
def __init__(self, min_momentum=None, min_momentum_ratio=None, **kwargs):
assert (min_momentum is None) ^ (min_momentum_ratio is None)
self.min_momentum = min_momentum
self.min_momentum_ratio = min_momentum_ratio
super(LinearAnnealingMomentumUpdaterHook, self).__init__(**kwargs)

def get_momentum(self, runner, base_momentum):
if self.by_epoch:
progress = runner.epoch
max_progress = runner.max_epochs
else:
progress = runner.iter
max_progress = runner.max_iters
if self.min_momentum_ratio is not None:
target_momentum = base_momentum * self.min_momentum_ratio
else:
target_momentum = self.min_momentum
return annealing_linear(base_momentum, target_momentum,
progress / max_progress)


@HOOKS.register_module()
class CyclicMomentumUpdaterHook(MomentumUpdaterHook):
"""Cyclic momentum Scheduler.
Expand Down
77 changes: 77 additions & 0 deletions tests/test_runner/test_hooks.py
Original file line number Diff line number Diff line change
Expand Up @@ -711,6 +711,83 @@ def test_cosine_runner_hook(multi_optimizers):
hook.writer.add_scalars.assert_has_calls(calls, any_order=True)


@pytest.mark.parametrize('multi_optimizers', (True, False))
def test_linear_runner_hook(multi_optimizers):
sys.modules['pavi'] = MagicMock()
loader = DataLoader(torch.ones((10, 2)))
runner = _build_demo_runner(multi_optimizers=multi_optimizers)

# add momentum scheduler

hook_cfg = dict(
type='LinearAnnealingMomentumUpdaterHook',
min_momentum_ratio=0.99 / 0.95,
by_epoch=False,
warmup_iters=2,
warmup_ratio=0.9 / 0.95)
runner.register_hook_from_cfg(hook_cfg)

# add momentum LR scheduler
hook_cfg = dict(
type='LinearAnnealingLrUpdaterHook',
by_epoch=False,
min_lr_ratio=0,
warmup_iters=2,
warmup_ratio=0.9)
runner.register_hook_from_cfg(hook_cfg)
runner.register_hook_from_cfg(dict(type='IterTimerHook'))
runner.register_hook(IterTimerHook())
# add pavi hook
hook = PaviLoggerHook(interval=1, add_graph=False, add_last_ckpt=True)
runner.register_hook(hook)
runner.run([loader], [('train', 1)])
shutil.rmtree(runner.work_dir)

# TODO: use a more elegant way to check values
assert hasattr(hook, 'writer')
if multi_optimizers:
calls = [
call(
'train', {
'learning_rate/model1': 0.02,
'learning_rate/model2': 0.01,
'momentum/model1': 0.95,
'momentum/model2': 0.9,
}, 1),
call(
'train', {
'learning_rate/model1': 0.01,
'learning_rate/model2': 0.005,
'momentum/model1': 0.97,
'momentum/model2': 0.9189473684210527,
}, 6),
call(
'train', {
'learning_rate/model1': 0.0019999999999999983,
'learning_rate/model2': 0.0009999999999999992,
'momentum/model1': 0.9860000000000001,
'momentum/model2': 0.9341052631578949,
}, 10)
]
else:
calls = [
call('train', {
'learning_rate': 0.02,
'momentum': 0.95
}, 1),
call('train', {
'learning_rate': 0.01,
'momentum': 0.97
}, 6),
call(
'train', {
'learning_rate': 0.0019999999999999983,
'momentum': 0.9860000000000001
}, 10)
]
hook.writer.add_scalars.assert_has_calls(calls, any_order=True)


@pytest.mark.parametrize('multi_optimizers, by_epoch', [(False, False),
(True, False),
(False, True),
Expand Down