mirrors
/
PaddleClas
mirror of https://github.com/PaddlePaddle/PaddleClas.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

Speedup EMA

pull/2747/head
Yang Nie 2023-03-02 03:40:25 +08:00 committed by Tingquan Gao
parent 541326eaf0
commit cabdc251fe
1 changed files with 7 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
ppcls/utils/ema.py
View File

@ -32,11 +32,14 @@ class ExponentialMovingAverage():
@paddle.no_grad() @paddle.no_grad()
def _update(self, model, update_fn): def _update(self, model, update_fn):
for ema_v, model_v in zip(self.module.state_dict().values(), model.state_dict().values()): for ema_v, model_v in zip(self.module.state_dict().values(),
ema_v.set_value(update_fn(ema_v, model_v)) model.state_dict().values()):
ema_v.set_value(update_fn(ema_v.numpy(), model_v.numpy()))
def update(self, model): def update(self, model):
self._update(model, update_fn=lambda e, m: self.decay * e + (1. - self.decay) * m) self._update(
model,
update_fn=lambda e, m: self.decay * e + (1. - self.decay) * m)
def set(self, model): def set(self, model):
self._update(model, update_fn=lambda e, m: m) self._update(model, update_fn=lambda e, m: m)