improve: implement logging metrics to tensorboard

Author: samedii

lantern/metric.py

@@ -2,25 +2,25 @@
 
 
 class ReduceMetric:
-    def __init__(self, reduce, compute, initial_state=None):
-        self._reduce = reduce
-        self._compute = compute
+    def __init__(self, reduce_fn, compute_fn, initial_state=None):
+        self.reduce_fn = reduce_fn
+        self.compute_fn = compute_fn
         self.state = initial_state
 
     def reduce(self, *args, **kwargs):
         return ReduceMetric(
-            reduce=self._reduce,
-            compute=self._compute,
-            initial_state=self._reduce(self.state, *args, **kwargs),
+            reduce_fn=self.reduce_fn,
+            compute_fn=self.compute_fn,
+            initial_state=self.reduce_fn(self.state, *args, **kwargs),
         )
 
     def compute(self):
-        return self._compute(self.state)
+        return self.compute_fn(self.state)
 
 
-def MapMetric(map, compute=np.mean):
+def MapMetric(map_fn, compute_fn=np.mean):
     return ReduceMetric(
-        reduce=lambda state, *args, **kwargs: state + [map(*args, **kwargs)],
-        compute=compute,
+        reduce_fn=lambda state, *args, **kwargs: state + [map_fn(*args, **kwargs)],
+        compute_fn=compute_fn,
         initial_state=list(),
     )

lantern/metrics.py

@@ -18,10 +18,11 @@ def compute(self):
 
 
 class Metrics:
-    def __init__(self, name, tensorboard_logger, metrics):
+    def __init__(self, name, tensorboard_logger, metrics, n_logs=0):
         self.name = name
         self.tensorboard_logger = tensorboard_logger
         self.metrics = metrics
+        self.n_logs = n_logs
 
     def __getitem__(self, name_or_names):
         if type(name_or_names) == str:
@@ -39,8 +40,17 @@ def update_(self, *args, **kwargs):
     def compute(self):
         return {name: metric.compute() for name, metric in self.metrics.items()}
 
-    def log_(self):
-        # TODO
+    def log_(self, step=None):
+        self.n_logs += 1
+        if step is None:
+            step = self.n_logs
+
+        for name, metric in self.metrics.items():
+            self.tensorboard_logger.add_scalar(
+                f"{self.name}/{name}",
+                metric.compute(),
+                step,
+            )
         return self
 
     def table(self):