Something went wrong. Try again.
This repository has no description
Something went wrong. Try again.
1.2 kB · 40 lines
Python
at main
1234567891011121314151617181920212223242526272829303132333435363738394041import copyimport torch.nn as nn
class EMAHelper(object): def __init__(self, mu=0.999): self.mu = mu self.shadow = {}
def register(self, module): if isinstance(module, nn.DataParallel): module = module.module for name, param in module.named_parameters(): if param.requires_grad: self.shadow[name] = param.data.clone()
def update(self, module): if isinstance(module, nn.DataParallel): module = module.module for name, param in module.named_parameters(): if param.requires_grad: self.shadow[name].data = (1. - self.mu) * param.data + self.mu * self.shadow[name].data
def ema(self, module): if isinstance(module, nn.DataParallel): module = module.module for name, param in module.named_parameters(): if param.requires_grad: param.data.copy_(self.shadow[name].data)
def ema_copy(self, module): module_copy = copy.deepcopy(module) self.ema(module_copy) return module_copy
def state_dict(self): return self.shadow
def load_state_dict(self, state_dict): self.shadow = state_dict