Speedup EMA

flytocc · Mar 1, 2023 · 35fc732 · 35fc732
1 parent 4da55fe
commit 35fc732
Showing 1 changed file with 7 additions and 4 deletions.
diff --git a/ppcls/utils/ema.py b/ppcls/utils/ema.py
@@ -32,11 +32,14 @@ def __init__(self, model, decay=0.9999):
 
     @paddle.no_grad()
     def _update(self, model, update_fn):
-        for ema_v, model_v in zip(self.module.state_dict().values(), model.state_dict().values()):
-            ema_v.set_value(update_fn(ema_v, model_v))
+        for ema_v, model_v in zip(self.module.state_dict().values(),
+                                  model.state_dict().values()):
+            ema_v.set_value(update_fn(ema_v.numpy(), model_v.numpy()))
 
     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):
-        self._update(model, update_fn=lambda e, m: m)
+        self._update(model, update_fn=lambda e, m: m)