| 4 | |
| 5 | |
| 6 | class EMAHelper(object): |
| 7 | def __init__(self, mu=0.999): |
| 8 | self.mu = mu |
| 9 | self.shadow = {} |
| 10 | |
| 11 | def register(self, module): |
| 12 | if isinstance(module, nn.DataParallel): |
| 13 | module = module.module |
| 14 | for name, param in module.named_parameters(): |
| 15 | if param.requires_grad: |
| 16 | self.shadow[name] = param.data.clone() |
| 17 | |
| 18 | def update(self, module): |
| 19 | if isinstance(module, nn.DataParallel): |
| 20 | module = module.module |
| 21 | for name, param in module.named_parameters(): |
| 22 | if param.requires_grad: |
| 23 | self.shadow[name].data = ( |
| 24 | 1. - |
| 25 | self.mu) * param.data + self.mu * self.shadow[name].data |
| 26 | |
| 27 | def ema(self, module): |
| 28 | if isinstance(module, nn.DataParallel): |
| 29 | module = module.module |
| 30 | for name, param in module.named_parameters(): |
| 31 | if param.requires_grad: |
| 32 | param.data.copy_(self.shadow[name].data) |
| 33 | |
| 34 | def ema_copy(self, module): |
| 35 | if isinstance(module, nn.DataParallel): |
| 36 | inner_module = module.module |
| 37 | module_copy = type(inner_module)(inner_module.config).to( |
| 38 | inner_module.config.device) |
| 39 | module_copy.load_state_dict(inner_module.state_dict()) |
| 40 | module_copy = nn.DataParallel(module_copy) |
| 41 | else: |
| 42 | module_copy = type(module)(module.config).to(module.config.device) |
| 43 | module_copy.load_state_dict(module.state_dict()) |
| 44 | # module_copy = copy.deepcopy(module) |
| 45 | self.ema(module_copy) |
| 46 | return module_copy |
| 47 | |
| 48 | def state_dict(self): |
| 49 | return self.shadow |
| 50 | |
| 51 | def load_state_dict(self, state_dict): |
| 52 | self.shadow = state_dict |