| import torch |
| import torch.nn as nn |
| from torch.nn.modules.batchnorm import _BatchNorm |
| import copy |
|
|
|
|
| class EMAModel(nn.Module): |
| def __init__( |
| self, |
| model, |
| update_after_step=0, |
| inv_gamma=1.0, |
| power=2 / 3, |
| min_value=0.0, |
| max_value=0.9999, |
| ): |
| super(EMAModel, self).__init__() |
| self.averaged_model = copy.deepcopy(model) |
| self.averaged_model.eval() |
| self.averaged_model.requires_grad_(False) |
|
|
| self.update_after_step = update_after_step |
| self.inv_gamma = inv_gamma |
| self.power = power |
| self.min_value = torch.tensor(min_value) |
| self.max_value = torch.tensor(max_value) |
|
|
| self.register_buffer("decay", torch.tensor(0.0)) |
| self.register_buffer("optimization_step", torch.tensor(0, dtype=torch.int64)) |
|
|
| def get_decay(self, optimization_step): |
| step = max(0, optimization_step - self.update_after_step - 1) |
| value = 1 - (1 + step / self.inv_gamma) ** -self.power |
| if step <= 0: |
| return torch.tensor(0.0) |
| return max(self.min_value, min(value, self.max_value)) |
|
|
| @torch.no_grad() |
| def step(self, new_model): |
| self.decay = self.get_decay(self.optimization_step).to(new_model.device) |
| for module, ema_module in zip( |
| new_model.modules(), self.averaged_model.modules() |
| ): |
| for param, ema_param in zip( |
| module.parameters(recurse=False), ema_module.parameters(recurse=False) |
| ): |
| if isinstance(module, _BatchNorm) or not param.requires_grad: |
| ema_param.copy_(param.to(dtype=ema_param.dtype).data) |
| else: |
| ema_param.mul_(self.decay) |
| ema_param.add_( |
| param.data.to(dtype=ema_param.dtype), alpha=1 - self.decay |
| ) |
| self.optimization_step += 1 |
|
|
