Implementation from https://fyubang.com/2019/06/01/ema/
| 344 | # self.model.load_state_dict(self.backup) |
| 345 | # self.backup = {} |
| 346 | class EMA: |
| 347 | """ |
| 348 | Implementation from https://fyubang.com/2019/06/01/ema/ |
| 349 | """ |
| 350 | |
| 351 | def __init__(self, model, decay): |
| 352 | self.model = model |
| 353 | self.decay = decay |
| 354 | self.shadow = {} |
| 355 | self.backup = {} |
| 356 | |
| 357 | def load(self, ema_model): |
| 358 | for name, param in ema_model.named_parameters(): |
| 359 | self.shadow[name] = param.data.clone() |
| 360 | |
| 361 | def register(self): |
| 362 | for name, param in self.model.named_parameters(): |
| 363 | if param.requires_grad: |
| 364 | self.shadow[name] = param.data.clone() |
| 365 | |
| 366 | def update(self): |
| 367 | for name, param in self.model.named_parameters(): |
| 368 | if param.requires_grad: |
| 369 | assert name in self.shadow |
| 370 | new_average = (1.0 - self.decay) * param.data + self.decay * self.shadow[name] |
| 371 | self.shadow[name] = new_average.clone() |
| 372 | |
| 373 | def apply_shadow(self): |
| 374 | for name, param in self.model.named_parameters(): |
| 375 | if param.requires_grad: |
| 376 | assert name in self.shadow |
| 377 | self.backup[name] = param.data |
| 378 | param.data = self.shadow[name] |
| 379 | |
| 380 | def restore(self): |
| 381 | for name, param in self.model.named_parameters(): |
| 382 | if param.requires_grad: |
| 383 | assert name in self.backup |
| 384 | param.data = self.backup[name] |
| 385 | self.backup = {} |
| 386 | |
| 387 | |
| 388 | class Bn_Controller: |