| 454 | |
| 455 | |
| 456 | class ModelEma(torch.nn.Module): |
| 457 | def __init__(self, model, decay=0.9997, device=None): |
| 458 | super(ModelEma, self).__init__() |
| 459 | # make a copy of the model for accumulating moving average of weights |
| 460 | self.module = deepcopy(model) |
| 461 | self.module.eval() |
| 462 | |
| 463 | # import pdb; pdb.set_trace() |
| 464 | |
| 465 | self.decay = decay |
| 466 | self.device = device # perform ema on different device from model if set |
| 467 | if self.device is not None: |
| 468 | self.module.to(device=device) |
| 469 | |
| 470 | def _update(self, model, update_fn): |
| 471 | with torch.no_grad(): |
| 472 | for ema_v, model_v in zip(self.module.state_dict().values(), |
| 473 | model.state_dict().values()): |
| 474 | if self.device is not None: |
| 475 | model_v = model_v.to(device=self.device) |
| 476 | ema_v.copy_(update_fn(ema_v, model_v)) |
| 477 | |
| 478 | def update(self, model): |
| 479 | self._update(model, |
| 480 | update_fn=lambda e, m: self.decay * e + |
| 481 | (1. - self.decay) * m) |
| 482 | |
| 483 | def set(self, model): |
| 484 | self._update(model, update_fn=lambda e, m: m) |
| 485 | |
| 486 | |
| 487 | class BestMetricSingle(): |