mirrors/pytorch-image-models
1
0
Fork 0
mirror of https://github.com/huggingface/pytorch-image-models.git synced 2025-06-03 15:01:08 +08:00
Code Issues Projects Releases Wiki Activity

Use in-place operations for EMA

Browse Source
This commit is contained in:
Jerome Rony 2022-11-17 11:53:29 -05:00
parent 25ffac6880
commit 6ec5cd6a99
1 changed files with 8 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
timm/utils/model_ema.py
View File

@ -117,10 +117,15 @@ class ModelEmaV2(nn.Module):
for ema_v, model_v in zip(self.module.state_dict().values(), model.state_dict().values()):
if self.device is not None:
model_v = model_v.to(device=self.device)
ema_v.copy_(update_fn(ema_v, model_v))
update_fn(ema_v, model_v)
def update(self, model):
self._update(model, update_fn=lambda e, m: self.decay * e + (1. - self.decay) * m)
def ema_update(e, m):
if m.is_floating_point():
e.mul_(self.decay).add_(m, alpha=1 - self.decay)
self._update(model, update_fn=ema_update)
def set(self, model):
self._update(model, update_fn=lambda e, m: m)
self._update(model, update_fn=lambda e, m: e.copy_(m))