1 files changed, 14 insertions, 8 deletions

diff --git a/training/util.py b/training/util.py
index 8bd8a83..61f1533 100644
--- a/training/util.py
+++ b/training/util.py
@@ -16,19 +16,25 @@ def save_args(basepath: Path, args, extra={}):
 
 
 class AverageMeter:
-    avg: Any
-
-    def __init__(self, name=None):
-        self.name = name
+    def __init__(self, inv_gamma=1.0, power=2 / 3):
+        self.inv_gamma = inv_gamma
+        self.power = power
         self.reset()
 
     def reset(self):
-        self.sum = self.count = self.avg = 0
+        self.step = 0
+        self.avg = 0
+
+    def get_decay(self):
+        if self.step <= 0:
+            return 1
+
+        return (self.step / self.inv_gamma) ** -self.power
 
     def update(self, val, n=1):
-        self.sum += val * n
-        self.count += n
-        self.avg = self.sum / self.count
+        for _ in range(n):
+            self.step += n
+            self.avg += (val - self.avg) * self.get_decay()
 
 
 class EMAModel(EMAModel_):