Newer
Older
# Copyright (c) OpenMMLab. All rights reserved.
import math
import warnings
import weakref
from collections import Counter
from functools import wraps
from typing import Callable, List, Optional, Union
from torch.optim import Optimizer
from mmengine.registry import PARAM_SCHEDULERS
INF = int(1e9)
OptimizerType = Union[OptimWrapper, Optimizer]
class _ParamScheduler:
"""Base class for parameter schedulers.
It should be inherited by all schedulers that schedule parameters in the
optimizer's ``param_groups``. All subclasses should overwrite the
``_get_value()`` according to their own schedule strategy.
The implementation is motivated by
https://github.com/pytorch/pytorch/blob/master/torch/optim/lr_scheduler.py.
Args:
optimizer (OptimWrapper or Optimizer): Wrapped optimizer.
param_name (str): Name of the parameter to be adjusted, such as
``lr``, ``momentum``.
begin (int): Step at which to start updating the parameters.
Defaults to 0.
end (int): Step at which to stop updating the parameters.
Defaults to INF.
last_step (int): The index of last step. Used for resuming without
state dict. Default value ``-1`` means the ``step`` function is
never be called before. Defaults to -1.
by_epoch (bool): Whether the scheduled parameters are updated by
epochs. Defaults to True.
verbose (bool): Whether to print the value for each update.
Defaults to False.
""" # noqa: E501
def __init__(self,
param_name: str,
begin: int = 0,
end: int = INF,
last_step: int = -1,
by_epoch: bool = True,
verbose: bool = False):
# Attach optimizer
if not isinstance(optimizer, (Optimizer, OptimWrapper)):
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
raise TypeError('``optimizer`` should be an Optimizer,'
'but got {}'.format(type(optimizer).__name__))
self.optimizer = optimizer
self.param_name = param_name
if end <= begin:
raise ValueError('end should be larger than begin, but got'
' begin={}, end={}'.format(begin, end))
self.begin = begin
self.end = end
self.by_epoch = by_epoch
assert isinstance(last_step, int) and last_step >= -1
# Initialize valid step count and base values
if last_step == -1:
for group in optimizer.param_groups:
# If the param is never be scheduled, record the current value
# as the initial value.
group.setdefault(f'initial_{param_name}', group[param_name])
else:
for i, group in enumerate(optimizer.param_groups):
if f'initial_{param_name}' not in group:
raise KeyError(
f"param 'initial_{param_name}' is not specified "
'in param_groups[{}] when resuming an optimizer'.
format(i))
self.base_values = [
group[f'initial_{param_name}'] for group in optimizer.param_groups
]
self.last_step = last_step
# Following https://github.com/pytorch/pytorch/issues/20124
# We would like to ensure that `scheduler.step()` is called after
# `optimizer.step()`
def with_counter(method: Callable):
if getattr(method, '_with_counter', False):
# `optimizer.step()` has already been replaced, return.
return method
# Keep a weak reference to the optimizer instance to prevent
# cyclic references.
instance_ref = weakref.ref(method.__self__) # type: ignore
# Get the unbound method for the same purpose.
func = method.__func__ # type: ignore
cls = instance_ref().__class__ # type: ignore
del method
@wraps(func)
def wrapper(*args, **kwargs):
instance = instance_ref()
instance._global_step += 1
wrapped = func.__get__(instance, cls)
return wrapped(*args, **kwargs)
# Note that the returned function here is no longer a bound method,
# so attributes like `__func__` and `__self__` no longer exist.
wrapper._with_counter = True # type: ignore
return wrapper
# add counter to optimizer
self.optimizer.step = with_counter(self.optimizer.step) # type: ignore
self.optimizer._global_step = -1 # type: ignore
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
self._global_step = -1
self.verbose = verbose
self.step()
def state_dict(self) -> dict:
"""Returns the state of the scheduler as a :class:`dict`.
It contains an entry for every variable in self.__dict__ which is not
the optimizer.
Returns:
dict: scheduler state.
"""
return {
key: value
for key, value in self.__dict__.items() if key != 'optimizer'
}
def load_state_dict(self, state_dict: dict):
"""Loads the schedulers state.
Args:
state_dict (dict): scheduler state. Should be an object returned
from a call to :meth:`state_dict`.
"""
self.__dict__.update(state_dict)
def get_last_value(self):
"""Return the last computed value by current scheduler.
Returns:
list: A list of the last computed value of the optimizer's
``param_group``.
"""
return self._last_value
def _get_value(self):
"""Compute value using chainable form of the scheduler."""
raise NotImplementedError
def print_value(self, is_verbose: bool, group: int, value: float):
"""Display the current parameter value.
Args:
is_verbose (bool): Whether to print the value.
group (int): The index of the current ``param_group``.
value (float): The parameter value.
"""
if is_verbose:
print('Adjusting parameter value'
' of group {} to {:.4e}.'.format(group, value))
def step(self):
"""Adjusts the parameter value of each parameter group based on the
specified schedule."""
# Raise a warning if old pattern is detected
# https://github.com/pytorch/pytorch/issues/20124
if self._global_step == 0:
if not hasattr(self.optimizer.step, '_with_counter'):
warnings.warn(
'Seems like `optimizer.step()` has been overridden after'
'parameter value scheduler initialization. Please, make'
'sure to call `optimizer.step()` before'
'`scheduler.step()`. See more details at'
'https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate', # noqa: E501
UserWarning)
# Just check if there were two first scheduler.step() calls
# before optimizer.step()
elif self.optimizer._global_step < 0:
warnings.warn(
'Detected call of `scheduler.step()` before'
'`optimizer.step()`. In PyTorch 1.1.0 and later, you'
'should call them in the opposite order: '
'`optimizer.step()` before `scheduler.step()`. '
'Failure to do this will result in PyTorch skipping '
'the first value of the parameter value schedule. '
'See more details at https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate', # noqa: E501
UserWarning)
self._global_step += 1
# Compute parameter value per param group in the effective range
if self.begin <= self._global_step < self.end:
self.last_step += 1
values = self._get_value()
for i, data in enumerate(zip(self.optimizer.param_groups, values)):
param_group, value = data
param_group[self.param_name] = value
self.print_value(self.verbose, i, value)
self._last_value = [
group[self.param_name] for group in self.optimizer.param_groups
]
@PARAM_SCHEDULERS.register_module()
class StepParamScheduler(_ParamScheduler):
"""Decays the parameter value of each parameter group by gamma every
step_size epochs. Notice that such decay can happen simultaneously with
other changes to the parameter value from outside this scheduler.
Args:
optimizer (OptimWrapper or Optimizer): Wrapped optimizer.
step_size (int): Period of parameter value decay.
gamma (float): Multiplicative factor of parameter value decay.
Defaults to 0.1.
begin (int): Step at which to start updating the parameters.
Defaults to 0.
end (int): Step at which to stop updating the parameters.
Defaults to INF.
last_step (int): The index of last step. Used for resume without
state dict. Defaults to -1.
by_epoch (bool): Whether the scheduled parameters are updated by
epochs. Defaults to True.
verbose (bool): Whether to print the value for each update.
Defaults to False.
"""
def __init__(self,
param_name: str,
step_size: int,
gamma: float = 0.1,
begin: int = 0,
end: int = INF,
last_step: int = -1,
by_epoch: bool = True,
verbose: bool = False):
self.step_size = step_size
self.gamma = gamma
super().__init__(
optimizer=optimizer,
param_name=param_name,
begin=begin,
end=end,
last_step=last_step,
by_epoch=by_epoch,
verbose=verbose)
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
@classmethod
def build_iter_from_epoch(cls,
*args,
step_size,
begin=0,
end=INF,
by_epoch=True,
epoch_length=None,
**kwargs):
"""Build an iter-based instance of this scheduler from an epoch-based
config."""
assert by_epoch, 'Only epoch-based kwargs whose `by_epoch=True` can ' \
'be converted to iter-based.'
assert epoch_length is not None and epoch_length > 0, \
f'`epoch_length` must be a positive integer, ' \
f'but got {epoch_length}.'
by_epoch = False
step_size = step_size * epoch_length
begin = begin * epoch_length
if end != INF:
end = end * epoch_length
return cls(
*args,
step_size=step_size,
begin=begin,
end=end,
by_epoch=by_epoch,
**kwargs)
"""Compute value using chainable form of the scheduler."""
if (self.last_step == 0) or (self.last_step % self.step_size != 0):
return [
group[self.param_name] for group in self.optimizer.param_groups
]
return [
group[self.param_name] * self.gamma
for group in self.optimizer.param_groups
]
@PARAM_SCHEDULERS.register_module()
class MultiStepParamScheduler(_ParamScheduler):
"""Decays the specified parameter in each parameter group by gamma once the
number of epoch reaches one of the milestones. Notice that such decay can
happen simultaneously with other changes to the parameter from outside this
scheduler.
Args:
optimizer (OptimWrapper or Optimizer): Wrapped optimizer.
milestones (list): List of epoch indices. Must be increasing.
gamma (float): Multiplicative factor of parameter value decay.
Defaults to 0.1.
begin (int): Step at which to start updating the parameters.
Defaults to 0.
end (int): Step at which to stop updating the parameters.
Defaults to INF.
last_step (int): The index of last step. Used for resume without
state dict. Defaults to -1.
by_epoch (bool): Whether the scheduled parameters are updated by
epochs. Defaults to True.
verbose (bool): Whether to print the value for each update.
Defaults to False.
"""
def __init__(self,
param_name: str,
milestones: List[int],
gamma: float = 0.1,
last_step: int = -1,
begin: int = 0,
end: int = INF,
by_epoch: bool = True,
verbose: bool = False):
self.milestones = Counter(milestones)
self.gamma = gamma
super().__init__(
optimizer,
param_name=param_name,
begin=begin,
end=end,
last_step=last_step,
by_epoch=by_epoch,
verbose=verbose)
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
@classmethod
def build_iter_from_epoch(cls,
*args,
milestones,
begin=0,
end=INF,
by_epoch=True,
epoch_length=None,
**kwargs):
"""Build an iter-based instance of this scheduler from an epoch-based
config."""
assert by_epoch, 'Only epoch-based kwargs whose `by_epoch=True` can ' \
'be converted to iter-based.'
assert epoch_length is not None and epoch_length > 0, \
f'`epoch_length` must be a positive integer, ' \
f'but got {epoch_length}.'
by_epoch = False
milestones = [i * epoch_length for i in milestones]
begin = begin * epoch_length
if end != INF:
end = end * epoch_length
return cls(
*args,
milestones=milestones,
begin=begin,
end=end,
by_epoch=by_epoch,
**kwargs)
"""Compute value using chainable form of the scheduler."""
if self.last_step not in self.milestones:
return [
group[self.param_name] for group in self.optimizer.param_groups
]
return [
group[self.param_name] *
self.gamma**self.milestones[self.last_step]
for group in self.optimizer.param_groups
]
@PARAM_SCHEDULERS.register_module()
class ConstantParamScheduler(_ParamScheduler):
"""Decays the parameter value of each parameter group by a small constant
factor until the number of epoch reaches a pre-defined milestone: ``end``.
Notice that such decay can happen simultaneously with other changes to the
parameter value from outside this scheduler.
Args:
optimizer (Optimizer or OptimWrapper): optimizer or Wrapped
optimizer.
factor (float): The number we multiply parameter value until the
milestone. Defaults to 1./3.
begin (int): Step at which to start updating the parameters.
Defaults to 0.
end (int): Step at which to stop updating the parameters.
Defaults to INF.
last_step (int): The index of last step. Used for resume without
state dict. Defaults to -1.
by_epoch (bool): Whether the scheduled parameters are updated by
epochs. Defaults to True.
verbose (bool): Whether to print the value for each update.
Defaults to False.
"""
def __init__(self,
param_name: str,
factor: float = 1.0 / 3,
begin: int = 0,
end: int = INF,
last_step: int = -1,
by_epoch: bool = True,
verbose: bool = False):
if factor > 1.0 or factor < 0:
raise ValueError(
'Constant multiplicative factor should between 0 and 1.')
self.factor = factor
self.total_iters = end - begin - 1
super().__init__(
optimizer,
param_name=param_name,
begin=begin,
end=end,
last_step=last_step,
by_epoch=by_epoch,
verbose=verbose)
@classmethod
def build_iter_from_epoch(cls,
*args,
begin=0,
end=INF,
by_epoch=True,
epoch_length=None,
**kwargs):
"""Build an iter-based instance of this scheduler from an epoch-based
config."""
assert by_epoch, 'Only epoch-based kwargs whose `by_epoch=True` can ' \
'be converted to iter-based.'
assert epoch_length is not None and epoch_length > 0, \
f'`epoch_length` must be a positive integer, ' \
f'but got {epoch_length}.'
by_epoch = False
begin = begin * epoch_length
if end != INF:
end = end * epoch_length
return cls(*args, begin=begin, end=end, by_epoch=by_epoch, **kwargs)
"""Compute value using chainable form of the scheduler."""
if self.last_step == 0:
return [
group[self.param_name] * self.factor
for group in self.optimizer.param_groups
]
if (self.last_step > self.total_iters
or (self.last_step != self.total_iters)):
return [
group[self.param_name] for group in self.optimizer.param_groups
]
if self.last_step == self.total_iters:
return [
group[self.param_name] * (1.0 / self.factor)
for group in self.optimizer.param_groups
]
@PARAM_SCHEDULERS.register_module()
class ExponentialParamScheduler(_ParamScheduler):
"""Decays the parameter value of each parameter group by gamma every epoch.
Args:
optimizer (Optimizer or OptimWrapper): optimizer or Wrapped
optimizer.
gamma (float): Multiplicative factor of parameter value decay.
begin (int): Step at which to start updating the parameters.
Defaults to 0.
end (int): Step at which to stop updating the parameters.
Defaults to INF.
last_step (int): The index of last step. Used for resume without
state dict. Defaults to -1.
by_epoch (bool): Whether the scheduled parameters are updated by
epochs. Defaults to True.
verbose (bool): Whether to print the value for each update.
Defaults to False.
"""
def __init__(self,
param_name: str,
gamma: float,
begin: int = 0,
end: int = INF,
last_step: int = -1,
by_epoch: bool = True,
verbose: bool = False):
self.gamma = gamma
super().__init__(
optimizer,
param_name=param_name,
begin=begin,
end=end,
last_step=last_step,
by_epoch=by_epoch,
verbose=verbose)
@classmethod
def build_iter_from_epoch(cls,
*args,
begin=0,
end=INF,
by_epoch=True,
epoch_length=None,
**kwargs):
"""Build an iter-based instance of this scheduler from an epoch-based
config."""
assert by_epoch, 'Only epoch-based kwargs whose `by_epoch=True` can ' \
'be converted to iter-based.'
assert epoch_length is not None and epoch_length > 0, \
f'`epoch_length` must be a positive integer, ' \
f'but got {epoch_length}.'
by_epoch = False
begin = begin * epoch_length
if end != INF:
end = end * epoch_length
return cls(*args, begin=begin, end=end, by_epoch=by_epoch, **kwargs)
"""Compute value using chainable form of the scheduler."""
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
if self.last_step == 0:
return [
group[self.param_name] for group in self.optimizer.param_groups
]
return [
group[self.param_name] * self.gamma
for group in self.optimizer.param_groups
]
@PARAM_SCHEDULERS.register_module()
class CosineAnnealingParamScheduler(_ParamScheduler):
r"""Set the parameter value of each parameter group using a cosine annealing
schedule, where :math:`\eta_{max}` is set to the initial value and
:math:`T_{cur}` is the number of epochs since the last restart in SGDR:
.. math::
\begin{aligned}
\eta_t & = \eta_{min} + \frac{1}{2}(\eta_{max} - \eta_{min})\left(1
+ \cos\left(\frac{T_{cur}}{T_{max}}\pi\right)\right),
& T_{cur} \neq (2k+1)T_{max}; \\
\eta_{t+1} & = \eta_{t} + \frac{1}{2}(\eta_{max} - \eta_{min})
\left(1 - \cos\left(\frac{1}{T_{max}}\pi\right)\right),
& T_{cur} = (2k+1)T_{max}.
\end{aligned}
Notice that because the schedule
is defined recursively, the parameter value can be simultaneously modified
outside this scheduler by other operators. If the parameter value is set
solely by this scheduler, the parameter value at each step becomes:
.. math::
\eta_t = \eta_{min} + \frac{1}{2}(\eta_{max} - \eta_{min})\left(1 +
\cos\left(\frac{T_{cur}}{T_{max}}\pi\right)\right)
It has been proposed in
`SGDR: Stochastic Gradient Descent with Warm Restarts`_. Note that this
only implements the cosine annealing part of SGDR, and not the restarts.
Args:
optimizer (Optimizer or OptimWrapper): optimizer or Wrapped
optimizer.
T_max (int): Maximum number of iterations.
eta_min (float): Minimum parameter value. Defaults to 0.
begin (int): Step at which to start updating the parameters.
Defaults to 0.
end (int): Step at which to stop updating the parameters.
Defaults to INF.
last_step (int): The index of last step. Used for resume without
state dict. Defaults to -1.
by_epoch (bool): Whether the scheduled parameters are updated by
epochs. Defaults to True.
verbose (bool): Whether to print the value for each update.
Defaults to False.
.. _SGDR\: Stochastic Gradient Descent with Warm Restarts:
https://arxiv.org/abs/1608.03983
"""
def __init__(self,
optimizer: Union[Optimizer, OptimWrapper],
param_name: str,
T_max: int,
eta_min: float = 0.,
begin: int = 0,
end: int = INF,
last_step: int = -1,
by_epoch: bool = True,
verbose: bool = False):
self.T_max = T_max
self.eta_min = eta_min
super().__init__(
optimizer,
param_name=param_name,
begin=begin,
end=end,
last_step=last_step,
by_epoch=by_epoch,
verbose=verbose)
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
@classmethod
def build_iter_from_epoch(cls,
*args,
T_max,
begin=0,
end=INF,
by_epoch=True,
epoch_length=None,
**kwargs):
"""Build an iter-based instance of this scheduler from an epoch-based
config."""
assert by_epoch, 'Only epoch-based kwargs whose `by_epoch=True` can ' \
'be converted to iter-based.'
assert epoch_length is not None and epoch_length > 0, \
f'`epoch_length` must be a positive integer, ' \
f'but got {epoch_length}.'
by_epoch = False
T_max = T_max * epoch_length
begin = begin * epoch_length
if end != INF:
end = end * epoch_length
return cls(
*args,
T_max=T_max,
begin=begin,
end=end,
by_epoch=by_epoch,
**kwargs)
"""Compute value using chainable form of the scheduler."""
if self.last_step == 0:
return [
group[self.param_name] for group in self.optimizer.param_groups
]
elif (self.last_step - 1 - self.T_max) % (2 * self.T_max) == 0:
return [
group[self.param_name] + (base_value - self.eta_min) *
(1 - math.cos(math.pi / self.T_max)) / 2
for base_value, group in zip(self.base_values,
self.optimizer.param_groups)
]
return [(1 + math.cos(math.pi * self.last_step / self.T_max)) /
(1 + math.cos(math.pi * (self.last_step - 1) / self.T_max)) *
(group[self.param_name] - self.eta_min) + self.eta_min
for group in self.optimizer.param_groups]
@PARAM_SCHEDULERS.register_module()
class LinearParamScheduler(_ParamScheduler):
"""Decays the parameter value of each parameter group by linearly changing
small multiplicative factor until the number of epoch reaches a pre-defined
milestone: ``end``.
Notice that such decay can happen simultaneously with other changes to the
parameter value from outside this scheduler.
optimizer (Optimizer or OptimWrapper): optimizer or Wrapped
optimizer.
start_factor (float): The number we multiply parameter value in the
first epoch. The multiplication factor changes towards end_factor
in the following epochs. Defaults to 1./3.
end_factor (float): The number we multiply parameter value at the end
of linear changing process. Defaults to 1.0.
begin (int): Step at which to start updating the parameters.
Defaults to 0.
end (int): Step at which to stop updating the parameters.
Defaults to INF.
last_step (int): The index of last step. Used for resume without
state dict. Defaults to -1.
by_epoch (bool): Whether the scheduled parameters are updated by
epochs. Defaults to True.
verbose (bool): Whether to print the value for each update.
Defaults to False.
"""
def __init__(self,
optimizer: Union[Optimizer, OptimWrapper],
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
param_name: str,
start_factor: float = 1.0 / 3,
end_factor: float = 1.0,
begin: int = 0,
end: int = INF,
last_step: int = -1,
by_epoch: bool = True,
verbose: bool = False):
if start_factor > 1.0 or start_factor < 0:
raise ValueError(
'Starting multiplicative factor should between 0 and 1.')
if end_factor > 1.0 or end_factor < 0:
raise ValueError(
'Ending multiplicative factor should between 0 and 1.')
self.start_factor = start_factor
self.end_factor = end_factor
self.total_iters = end - begin - 1
super().__init__(
optimizer,
param_name=param_name,
begin=begin,
end=end,
last_step=last_step,
by_epoch=by_epoch,
verbose=verbose)
@classmethod
def build_iter_from_epoch(cls,
*args,
begin=0,
end=INF,
by_epoch=True,
epoch_length=None,
**kwargs):
"""Build an iter-based instance of this scheduler from an epoch-based
config."""
assert by_epoch, 'Only epoch-based kwargs whose `by_epoch=True` can ' \
'be converted to iter-based.'
assert epoch_length is not None and epoch_length > 0, \
f'`epoch_length` must be a positive integer, ' \
f'but got {epoch_length}.'
by_epoch = False
begin = begin * epoch_length
if end != INF:
end = end * epoch_length
return cls(*args, begin=begin, end=end, by_epoch=by_epoch, **kwargs)
"""Compute value using chainable form of the scheduler."""
if self.last_step == 0:
return [
group[self.param_name] * self.start_factor
for group in self.optimizer.param_groups
]
return [
group[self.param_name] *
(1. + (self.end_factor - self.start_factor) /
(self.total_iters * self.start_factor + (self.last_step - 1) *
(self.end_factor - self.start_factor)))
for group in self.optimizer.param_groups
]
@PARAM_SCHEDULERS.register_module()
class PolyParamScheduler(_ParamScheduler):
"""Decays the parameter value of each parameter group in a polynomial decay
scheme.
Notice that such decay can happen simultaneously with other changes to the
parameter value from outside this scheduler.
Args:
optimizer (Optimizer or OptimWrapper): optimizer or Wrapped
optimizer.
eta_min (float): Minimum parameter value at the end of scheduling.
Defaults to 0.
power (float): The power of the polynomial. Defaults to 1.0.
begin (int): Step at which to start updating the parameters.
Defaults to 0.
end (int): Step at which to stop updating the parameters.
Defaults to INF.
last_step (int): The index of last step. Used for resume without
state dict. Defaults to -1.
by_epoch (bool): Whether the scheduled parameters are updated by
epochs. Defaults to True.
verbose (bool): Whether to print the value for each update.
Defaults to False.
"""
def __init__(self,
optimizer: Union[Optimizer, OptimWrapper],
param_name: str,
eta_min: float = 0,
power: float = 1.0,
begin: int = 0,
end: int = INF,
last_step: int = -1,
by_epoch: bool = True,
verbose: bool = False):
self.eta_min = eta_min
self.power = power
self.total_iters = end - begin - 1
super().__init__(
optimizer,
param_name=param_name,
begin=begin,
end=end,
last_step=last_step,
by_epoch=by_epoch,
verbose=verbose)
@classmethod
def build_iter_from_epoch(cls,
*args,
begin=0,
end=INF,
by_epoch=True,
epoch_length=None,
**kwargs):
"""Build an iter-based instance of this scheduler from an epoch-based
config."""
assert by_epoch, 'Only epoch-based kwargs whose `by_epoch=True` can ' \
'be converted to iter-based.'
assert epoch_length is not None and epoch_length > 0, \
f'`epoch_length` must be a positive integer, ' \
f'but got {epoch_length}.'
by_epoch = False
begin = begin * epoch_length
if end != INF:
end = end * epoch_length
return cls(*args, begin=begin, end=end, by_epoch=by_epoch, **kwargs)
def _get_value(self):
"""Compute value using chainable form of the scheduler."""
if self.last_step == 0:
return [
group[self.param_name] for group in self.optimizer.param_groups
]
return [(group[self.param_name] - self.eta_min) *
(1 - 1 / (self.total_iters - self.last_step + 1))**self.power +
self.eta_min for group in self.optimizer.param_groups]
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
@PARAM_SCHEDULERS.register_module()
class OneCycleParamScheduler(_ParamScheduler):
r"""Sets the parameters of each parameter group according to the
1cycle learning rate policy. The 1cycle policy anneals the learning
rate from an initial learning rate to some maximum learning rate and then
from that maximum learning rate to some minimum learning rate much lower
than the initial learning rate.
This policy was initially described in the paper `Super-Convergence:
Very Fast Training of Neural Networks Using Large Learning Rates`_.
The 1cycle learning rate policy changes the learning rate after every
batch. `step` should be called after a batch has been used for training.
This scheduler is not chainable.
Note also that the total number of steps in the cycle can be determined in
one of two ways (listed in order of precedence):
#. A value for total_steps is explicitly provided.
#. If total_steps is not defined, begin and end of the ParamSchedul will
works for it. In this case, the number of total steps is inferred by
total_steps = end - begin
The default behaviour of this scheduler follows the fastai implementation
of 1cycle, which claims that "unpublished work has shown even better
results by using only two phases". To mimic the behaviour of the original
paper instead, set ``three_phase=True``.
Args:
optimizer (Optimizer): Wrapped optimizer.
eta_max (float or list): Upper parameter value boundaries in the cycle
for each parameter group.
total_steps (int): The total number of steps in the cycle. Note that
if a value is not provided here, then it will be equal to
``end - begin``. Default to None
pct_start (float): The percentage of the cycle (in number of steps)
spent increasing the learning rate.
Default to 0.3
anneal_strategy (str): {'cos', 'linear'}
Specifies the annealing strategy: "cos" for cosine annealing,
"linear" for linear annealing.
Default to 'cos'
div_factor (float): Determines the initial learning rate via
initial_param = eta_max/div_factor
Default to 25
final_div_factor (float): Determines the minimum learning rate via
eta_min = initial_param/final_div_factor
Default to 1e4
three_phase (bool): If ``True``, use a third phase of the schedule to
annihilate the learning rate according to 'final_div_factor'
instead of modifying the second phase (the first two phases will be
symmetrical about the step indicated by 'pct_start').
last_step (int): The index of last step. Used for resume without
state dict. Defaults to -1.
by_epoch (bool): Whether the scheduled parameters are updated by
epochs. Defaults to True.
verbose (bool): Whether to print the value for each update.
Defaults to False.
.. _Super-Convergence\: Very Fast Training of Neural Networks Using Large Learning Rates:
https://arxiv.org/abs/1708.07120
"""# noqa E501
def __init__(self,
optimizer: Union[Optimizer, OptimWrapper],
param_name: str,
eta_max: float = 0,
total_steps: Optional[int] = None,
pct_start: float = 0.3,
anneal_strategy: str = 'cos',
div_factor: float = 25.,
final_div_factor: float = 1e4,
three_phase: bool = False,
begin: int = 0,
end: int = INF,
last_step: int = -1,
by_epoch: bool = True,
verbose: bool = False):
assert param_name == 'lr', ('OneCycle only works for learning rate '
'updating, but got patam_name as '
f'{param_name}')
self.eta_max = eta_max
self.div_factor = div_factor
self.final_div_factor = final_div_factor
# Validate total_steps
if total_steps is not None:
if total_steps <= 0 or not isinstance(total_steps, int):
raise ValueError('Expected positive integer total_steps, '
f'but got {total_steps}')
self.total_steps = total_steps
else:
self.total_steps = self.end - self.begin
# Validate pct_start
if pct_start < 0 or pct_start > 1 or not isinstance(pct_start, float):
raise ValueError('Expected float between 0 and 1 pct_start, '
f'but got {pct_start}')
# Validate anneal_strategy
if anneal_strategy not in ['cos', 'linear']:
raise ValueError(
'anneal_strategy must by one of "cos" or "linear", '
f'instead got {anneal_strategy}')
elif anneal_strategy == 'cos':
self.anneal_func = self._annealing_cos
elif anneal_strategy == 'linear':
self.anneal_func = self._annealing_linear
if three_phase:
self._schedule_phases = [
{
'end_step': float(pct_start * self.total_steps) - 1,
f'start_{param_name}': f'initial_{param_name}',
f'end_{param_name}': f'max_{param_name}'
},
{
'end_step': float(2 * pct_start * self.total_steps) - 2,
f'start_{param_name}': f'max_{param_name}',
f'end_{param_name}': f'initial_{param_name}'
},
{
'end_step': self.total_steps - 1,
f'start_{param_name}': f'initial_{param_name}',
f'end_{param_name}': f'min_{param_name}'
},
]
else:
self._schedule_phases = [
{
'end_step': float(pct_start * self.total_steps) - 1,
f'start_{param_name}': f'initial_{param_name}',
f'end_{param_name}': f'max_{param_name}'
},
{
'end_step': self.total_steps - 1,
f'start_{param_name}': f'max_{param_name}',
f'end_{param_name}': f'min_{param_name}'
},
]
# Initialize parameters
max_values = self._format_param(f'max_{param_name}', optimizer,
eta_max)
if last_step == -1:
for idx, group in enumerate(optimizer.param_groups):
group[f'initial_{param_name}'] = max_values[idx] / div_factor
group[f'max_{param_name}'] = max_values[idx]
group[f'min_{param_name}'] = \
group[f'initial_{param_name}'] / final_div_factor