RateTracker smooth every time we report instead of every time we update

torch_xla_py/xla_model.py

@@ -144,28 +144,40 @@ def __init__(self, smooth_factor=0.8):
     self._smooth_factor = smooth_factor
     self._start_time = time.time()
     self._partial_time = self._start_time
-    self._count = 0
-    self._rate = 0.0
+    self._partial_count = 0.0
+    self._partial_rate = None
+    self._count = 0.0
 
-  def update(self, count):
-    now = time.time()
-    delta = now - self._partial_time
-    if delta > 0:
-      rate = (count - self._count) / delta
-      self._rate = (
-          self._rate * self._smooth_factor + rate * (1.0 - self._smooth_factor))
+  def _update(self, now, rate):
+    self._partial_count += self._count
+    self._count = 0.0
     self._partial_time = now
-    self._count = count
-    return self._rate
+    self._partial_rate = rate
 
   def add(self, count):
-    return self.update(self._count + count)
+    self._count += count
+
+  def _smooth(self, current_rate):
+    if self._partial_rate is None:
+      smoothed_rate = current_rate
+    else:
+      smoothed_rate = ((1 - self._smooth_factor) * current_rate +
+                       self._smooth_factor * self._partial_rate)
+    return smoothed_rate
 
   def rate(self):
-    return self._rate
+    now = time.time()
+    delta = now - self._partial_time
+    report_rate = 0.0
+    if delta > 0:
+      report_rate = self._smooth(self._count / delta)
+      self._update(now, report_rate)
+    return report_rate
 
   def global_rate(self):
-    return self._count / (self._partial_time - self._start_time)
+    delta = time.time() - self._start_time
+    count = self._partial_count + self._count
+    return count / delta if delta > 0 else 0.0
 
 
 class TrainStepMetrics(object):