improve: handle single value and dict metrics

Author: samedii

lantern/metric.py

@@ -5,54 +5,42 @@
 from pydantic import BaseModel, Extra
 
 
-class MapMetric(BaseModel):
-    map_fn_: Optional[Callable[..., Any]]
-    # map_fn: Optional[Callable] = lambda value: self, value  # HACK: why are we getting self?
+class MapMetric(FunctionalBase):
+    map_fn: Optional[Callable[..., Any]]
     state: List[Any]
 
     class Config:
         arbitrary_types_allowed = True
         allow_mutation = True
         extra = Extra.forbid
 
-    def __init__(self, map_fn_=None, state=list()):
+    def __init__(self, map_fn=None, state=list()):
         super().__init__(
-            map_fn_=map_fn_,
+            map_fn=map_fn,
             state=state,
         )
 
-    def replace(self, **kwargs):
-        new_dict = self.dict()
-        new_dict.update(**kwargs)
-        return type(self)(**new_dict)
-
     def map(self, fn):
-        # return self.replace(fn=lambda value: fn(self.map_fn_(value)))
-        # HACK: why doesn't the above work?
-        if self.map_fn_ is None:
-            return MapMetric(
-                map_fn_=fn,
-                state=self.state,
-            )
+        if self.map_fn is None:
+            return self.replace(map_fn=fn)
         else:
-            return MapMetric(
-                map_fn_=lambda *args, **kwargs: fn(self.map_fn_(*args, **kwargs)),
-                state=self.state,
+            return self.replace(
+                map_fn=lambda *args, **kwargs: fn(self.map_fn(*args, **kwargs))
             )
 
     def starmap(self, fn):
         return self.map(star(fn))
 
     def reduce(self, fn):
-        if self.map_fn_ is None:
+        if self.map_fn is None:
             return ReduceMetric(
-                map_fn_=lambda *args: args,
+                map_fn=lambda *args: args,
                 reduce_fn=lambda state, args: fn(state, *args),
                 state=self.state,  # TODO: apply function on state...
             )
         else:
             return ReduceMetric(
-                map_fn_=self.map_fn_,
+                map_fn=self.map_fn,
                 reduce_fn=fn,
                 state=self.state,
             )
@@ -64,22 +52,32 @@ def staraggregate(self, fn):
         return self.aggregate(star(fn))
 
     def update_(self, *args, **kwargs):
-        if self.map_fn_ is None:
+        if self.map_fn is None:
             self.state.append(args)
         else:
-            self.state.append(self.map_fn_(*args, **kwargs))
+            self.state.append(self.map_fn(*args, **kwargs))
         return self
 
     def update(self, *args, **kwargs):
-        if self.map_fn_ is None:
-            return self.replace(state=self.state + ([args[0]] if len(args) == 1 else [args]))
+        if self.map_fn is None:
+            return self.replace(
+                state=self.state + ([args[0]] if len(args) == 1 else [args])
+            )
         else:
-            return self.replace(state=self.state + [self.map_fn_(*args, **kwargs)])
+            return self.replace(state=self.state + [self.map_fn(*args, **kwargs)])
 
     def compute(self):
         return self.state
 
-    def log(self, tensorboard_logger, tag, step=None):
+    def log(self, tensorboard_logger, tag, metric_name, step=None):
+        tensorboard_logger.add_scalar(
+            f"{tag}/{metric_name}",
+            self.compute(),
+            step,
+        )
+        return self
+
+    def log_dict(self, tensorboard_logger, tag, step=None):
         for name, value in self.compute().items():
             tensorboard_logger.add_scalar(
                 f"{tag}/{name}",
@@ -92,34 +90,41 @@ def log(self, tensorboard_logger, tag, step=None):
 Metric = MapMetric
 
 
-class ReduceMetric(BaseModel):
-    map_fn_: Callable[..., Any]
+class ReduceMetric(FunctionalBase):
+    map_fn: Callable[..., Any]
     reduce_fn: Callable[..., Any]
     state: Any
 
     class Config:
         arbitrary_types_allowed = True
         allow_mutation = True
-        extra = Extra.forbid
 
     def replace(self, **kwargs):
         new_dict = self.dict()
         new_dict.update(**kwargs)
         return type(self)(**new_dict)
 
     def update_(self, *args, **kwargs):
-        self.state = self.reduce_fn(self.state, self.map_fn_(*args, **kwargs))
+        self.state = self.reduce_fn(self.state, self.map_fn(*args, **kwargs))
         return self
 
     def update(self, *args, **kwargs):
         return self.replace(
-            state=self.reduce_fn(self.state, self.map_fn_(*args, **kwargs))
+            state=self.reduce_fn(self.state, self.map_fn(*args, **kwargs))
         )
 
     def compute(self):
         return self.state
 
-    def log(self, tensorboard_logger, tag, step=None):
+    def log(self, tensorboard_logger, tag, metric_name, step=None):
+        tensorboard_logger.add_scalar(
+            f"{tag}/{metric_name}",
+            self.compute(),
+            step,
+        )
+        return self
+
+    def log_dict(self, tensorboard_logger, tag, step=None):
         for name, value in self.compute().items():
             tensorboard_logger.add_scalar(
                 f"{tag}/{name}",
@@ -129,24 +134,16 @@ def log(self, tensorboard_logger, tag, step=None):
         return self
 
 
-class AggregateMetric(BaseModel):
+class AggregateMetric(FunctionalBase):
     metric: Union[MapMetric, ReduceMetric]
     aggregate_fn: Callable
 
     class Config:
         arbitrary_types_allowed = True
         allow_mutation = True
-        extra = Extra.forbid
-
-    def replace(self, **kwargs):
-        new_dict = self.dict()
-        new_dict.update(**kwargs)
-        return type(self)(**new_dict)
 
     def map(self, fn):
-        return self.replace(
-            aggregate_fn=lambda state: fn(self.aggregate_fn(state))
-        )
+        return self.replace(aggregate_fn=lambda state: fn(self.aggregate_fn(state)))
 
     def starmap(self, fn):
         return self.map(star(fn))
@@ -161,7 +158,15 @@ def update(self, *args, **kwargs):
     def compute(self):
         return self.aggregate_fn(self.metric.compute())
 
-    def log(self, tensorboard_logger, tag, step=None):
+    def log(self, tensorboard_logger, tag, metric_name, step=None):
+        tensorboard_logger.add_scalar(
+            f"{tag}/{metric_name}",
+            self.compute(),
+            step,
+        )
+        return self
+
+    def log_dict(self, tensorboard_logger, tag, step=None):
         for name, value in self.compute().items():
             tensorboard_logger.add_scalar(
                 f"{tag}/{name}",

lantern/metric_table.py

@@ -21,11 +21,14 @@ def __init__(self, name, metrics):
         )
 
     def compute(self):
-        return {
-            metric_name: value
-            for metrics in self.metrics.values()
-            for metric_name, value in metrics.compute().items()
-        }
+        log_dict = dict()
+        for metric_name, metric in self.metrics.items():
+            metric_value = metric.compute()
+            if isinstance(metric_value, dict):
+                log_dict.update(**metric_value)
+            else:
+                log_dict[metric_name] = metric_value
+        return log_dict
 
     def table(self):
         return "\n".join(

lantern/progress_bar.py

@@ -12,10 +12,12 @@ def ProgressBar(data_loader, name, metrics: Optional[Dict[str, Metric]] = None):
         with tqdm(data_loader, desc=name, leave=False) as tqdm_:
             for item in tqdm_:
                 yield item
-                tqdm_.set_postfix(
-                    {
-                        name: value
-                        for metrics in metrics.values()
-                        for name, value in metrics.compute().items()
-                    }
-                )
+
+                log_dict = dict()
+                for metric_name, metric in metrics.items():
+                    metric_value = metric.compute()
+                    if isinstance(metric_value, dict):
+                        log_dict.update(**metric_value)
+                    else:
+                        log_dict[metric_name] = metric_value
+                tqdm_.set_postfix(log_dict)

test/test_mnist.py

@@ -55,7 +55,7 @@ def test_mnist():
     tensorboard_logger = torch.utils.tensorboard.SummaryWriter()
     early_stopping = lantern.EarlyStopping(tensorboard_logger=tensorboard_logger)
     train_metrics = dict(
-        loss=lantern.Metric().reduce(lambda state, loss: dict(loss=loss.item())),
+        loss=lantern.Metric().reduce(lambda state, loss: loss.item()),
     )
 
     for epoch in lantern.Epochs(2):
@@ -73,14 +73,16 @@ def test_mnist():
             train_metrics["loss"].update_(loss)
             sleep(0.5)
 
-            for metrics in train_metrics.values():
-                metrics.log(tensorboard_logger, "train", epoch)
+            for metric_name, metric in train_metrics.items():
+                metric.log(tensorboard_logger, "train", metric_name, epoch)
 
         print(lantern.MetricTable("train", train_metrics))
 
         evaluate_metrics = {
             name: dict(
-                loss=lantern.Metric().reduce(lambda state, loss: dict(loss=loss.item())),
+                loss=lantern.Metric().reduce(
+                    lambda state, loss: dict(loss=loss.item())
+                ),
             )
             for name in evaluate_data_loaders
         }
@@ -94,7 +96,7 @@ def test_mnist():
                 evaluate_metrics[name]["loss"].update_(loss)
 
             for metrics in evaluate_metrics[name].values():
-                metrics.log(tensorboard_logger, name, epoch)
+                metrics.log_dict(tensorboard_logger, name, epoch)
 
             print(lantern.MetricTable(name, evaluate_metrics[name]))