File size: 1,923 Bytes
a3f6fc6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 | 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
|