|
|
|
from ...parallel import is_module_wrapper
|
|
from ..hooks.hook import HOOKS, Hook
|
|
|
|
|
|
@HOOKS.register_module()
|
|
class EMAHook(Hook):
|
|
r"""Exponential Moving Average Hook.
|
|
|
|
Use Exponential Moving Average on all parameters of model in training
|
|
process. All parameters have a ema backup, which update by the formula
|
|
as below. EMAHook takes priority over EvalHook and CheckpointSaverHook.
|
|
|
|
.. math::
|
|
|
|
\text{Xema\_{t+1}} = (1 - \text{momentum}) \times
|
|
\text{Xema\_{t}} + \text{momentum} \times X_t
|
|
|
|
Args:
|
|
momentum (float): The momentum used for updating ema parameter.
|
|
Defaults to 0.0002.
|
|
interval (int): Update ema parameter every interval iteration.
|
|
Defaults to 1.
|
|
warm_up (int): During first warm_up steps, we may use smaller momentum
|
|
to update ema parameters more slowly. Defaults to 100.
|
|
resume_from (str): The checkpoint path. Defaults to None.
|
|
"""
|
|
|
|
def __init__(self,
|
|
momentum=0.0002,
|
|
interval=1,
|
|
warm_up=100,
|
|
resume_from=None):
|
|
assert isinstance(interval, int) and interval > 0
|
|
self.warm_up = warm_up
|
|
self.interval = interval
|
|
assert momentum > 0 and momentum < 1
|
|
self.momentum = momentum**interval
|
|
self.checkpoint = resume_from
|
|
|
|
def before_run(self, runner):
|
|
"""To resume model with it's ema parameters more friendly.
|
|
|
|
Register ema parameter as ``named_buffer`` to model
|
|
"""
|
|
model = runner.model
|
|
if is_module_wrapper(model):
|
|
model = model.module
|
|
self.param_ema_buffer = {}
|
|
self.model_parameters = dict(model.named_parameters(recurse=True))
|
|
for name, value in self.model_parameters.items():
|
|
|
|
buffer_name = f"ema_{name.replace('.', '_')}"
|
|
self.param_ema_buffer[name] = buffer_name
|
|
model.register_buffer(buffer_name, value.data.clone())
|
|
self.model_buffers = dict(model.named_buffers(recurse=True))
|
|
if self.checkpoint is not None:
|
|
runner.resume(self.checkpoint)
|
|
|
|
def after_train_iter(self, runner):
|
|
"""Update ema parameter every self.interval iterations."""
|
|
curr_step = runner.iter
|
|
|
|
momentum = min(self.momentum,
|
|
(1 + curr_step) / (self.warm_up + curr_step))
|
|
if curr_step % self.interval != 0:
|
|
return
|
|
for name, parameter in self.model_parameters.items():
|
|
buffer_name = self.param_ema_buffer[name]
|
|
buffer_parameter = self.model_buffers[buffer_name]
|
|
buffer_parameter.mul_(1 - momentum).add_(momentum, parameter.data)
|
|
|
|
def after_train_epoch(self, runner):
|
|
"""We load parameter values from ema backup to model before the
|
|
EvalHook."""
|
|
self._swap_ema_parameters()
|
|
|
|
def before_train_epoch(self, runner):
|
|
"""We recover model's parameter from ema backup after last epoch's
|
|
EvalHook."""
|
|
self._swap_ema_parameters()
|
|
|
|
def _swap_ema_parameters(self):
|
|
"""Swap the parameter of model with parameter in ema_buffer."""
|
|
for name, value in self.model_parameters.items():
|
|
temp = value.data.clone()
|
|
ema_buffer = self.model_buffers[self.param_ema_buffer[name]]
|
|
value.data.copy_(ema_buffer.data)
|
|
ema_buffer.data.copy_(temp)
|
|
|
