From 6d46bf79bd7710cea799fbfe27c12d06d12cd53f Mon Sep 17 00:00:00 2001
From: Volpeon <git@volpeon.ink>
Date: Thu, 27 Apr 2023 07:47:59 +0200
Subject: Update

---
 training/util.py | 22 ++++++++++++++--------
 1 file changed, 14 insertions(+), 8 deletions(-)

(limited to 'training/util.py')

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_):
-- 
cgit v1.2.3-54-g00ecf