|
import torch |
|
import torch.nn as nn |
|
from copy import deepcopy |
|
|
|
|
|
class EMA(nn.Module): |
|
def __init__(self, src_model: nn.Module, beta: float, copy: bool = True): |
|
super().__init__() |
|
if copy: |
|
self.model = deepcopy(src_model) |
|
else: |
|
self.model = src_model |
|
self.model.eval() |
|
self.model.requires_grad_(False) |
|
self.beta = beta |
|
|
|
def step(self, src_model): |
|
one_minus_beta = 1.0 - self.beta |
|
for ema_param, src_param in zip( |
|
self.model.parameters(), src_model.parameters() |
|
): |
|
|
|
ema_param.data.mul_(self.beta).add_(src_param.data, alpha=one_minus_beta) |
|
ema_param.requires_grad_(False) |
|
|
|
def forward(self, *args, **kwargs): |
|
with torch.no_grad(): |
|
return self.model(*args, **kwargs) |
|
|
