diff options
Diffstat (limited to 'torch_ema/ema.py')
| -rw-r--r-- | torch_ema/ema.py | 50 |
1 files changed, 31 insertions, 19 deletions
diff --git a/torch_ema/ema.py b/torch_ema/ema.py index 7771ef7..0233c78 100644 --- a/torch_ema/ema.py +++ b/torch_ema/ema.py @@ -1,23 +1,30 @@ from __future__ import division from __future__ import unicode_literals +from typing import Iterable + import torch -# Partially based on: https://github.com/tensorflow/tensorflow/blob/r1.13/tensorflow/python/training/moving_averages.py +# Partially based on: +# https://github.com/tensorflow/tensorflow/blob/r1.13/tensorflow/python/training/moving_averages.py class ExponentialMovingAverage: """ Maintains (exponential) moving average of a set of parameters. - """ - def __init__(self, parameters, decay, use_num_updates=True): - """ - Args: - parameters: Iterable of `torch.nn.Parameter`; usually the result of + + Args: + parameters: Iterable of `torch.nn.Parameter`; usually the result of `model.parameters()`. - decay: The exponential decay. - use_num_updates: Whether to use number of updates when computing + decay: The exponential decay. + use_num_updates: Whether to use number of updates when computing averages. - """ + """ + def __init__( + self, + parameters: Iterable[torch.nn.Parameter], + decay: float, + use_num_updates: bool = True + ): if decay < 0.0 or decay > 1.0: raise ValueError('Decay must be between 0 and 1') self.decay = decay @@ -26,7 +33,7 @@ class ExponentialMovingAverage: for p in parameters if p.requires_grad] self.collected_params = [] - def update(self, parameters): + def update(self, parameters: Iterable[torch.nn.Parameter]) -> None: """ Update currently maintained parameters. @@ -40,18 +47,24 @@ class ExponentialMovingAverage: decay = self.decay if self.num_updates is not None: self.num_updates += 1 - decay = min(decay, (1 + self.num_updates) / (10 + self.num_updates)) + decay = min( + decay, + (1 + self.num_updates) / (10 + self.num_updates) + ) one_minus_decay = 1.0 - decay with torch.no_grad(): parameters = [p for p in parameters if p.requires_grad] for s_param, param in zip(self.shadow_params, parameters): - s_param.sub_(one_minus_decay * (s_param - param)) + tmp = (s_param - param) + # tmp will be a new tensor so we can do in-place + tmp.mul_(one_minus_decay) + s_param.sub_(tmp) - def copy_to(self, parameters): + def copy_to(self, parameters: Iterable[torch.nn.Parameter]) -> None: """ Copy current parameters into given collection of parameters. - Args: + Args: parameters: Iterable of `torch.nn.Parameter`; the parameters to be updated with the stored moving averages. """ @@ -59,11 +72,11 @@ class ExponentialMovingAverage: if param.requires_grad: param.data.copy_(s_param.data) - def store(self, parameters): + def store(self, parameters: Iterable[torch.nn.Parameter]) -> None: """ Save the current parameters for restoring later. - Args: + Args: parameters: Iterable of `torch.nn.Parameter`; the parameters to be temporarily stored. """ @@ -71,7 +84,7 @@ class ExponentialMovingAverage: for param in parameters if param.requires_grad] - def restore(self, parameters): + def restore(self, parameters: Iterable[torch.nn.Parameter]) -> None: """ Restore the parameters stored with the `store` method. Useful to validate the model with EMA parameters without affecting the @@ -79,11 +92,10 @@ class ExponentialMovingAverage: `copy_to` method. After validation (or model saving), use this to restore the former parameters. - Args: + Args: parameters: Iterable of `torch.nn.Parameter`; the parameters to be updated with the stored parameters. """ for c_param, param in zip(self.collected_params, parameters): if param.requires_grad: param.data.copy_(c_param.data) - |