feature!: chain metrics

BREAKING CHANGE

Author: samedii

lantern/__init__.py

@@ -6,7 +6,7 @@
 from lantern.functional_base import FunctionalBase
 from lantern.tensor import Tensor
 from lantern.epochs import Epochs
-from lantern.metric import Metric, ReduceMetric, MapMetric
+from lantern.metric import Metric
 from lantern.metric_table import MetricTable
 from lantern.module_device import module_device
 from lantern.module_compose import ModuleCompose

lantern/metric.py

@@ -1,61 +1,175 @@
 import numpy as np
-from abc import ABC, abstractmethod
 from lantern import FunctionalBase
-from typing import Callable, Any, Optional
+from lantern.functional import star
+from typing import Callable, Any, Optional, Dict, List, Union
+from pydantic import BaseModel, Extra
 
 
-class Metric(ABC):
-    @abstractmethod
-    def update(self):
-        ...
+class MapMetric(BaseModel):
+    map_fn_: Optional[Callable[..., Any]]
+    # map_fn: Optional[Callable] = lambda value: self, value  # HACK: why are we getting self?
+    state: List[Any]
 
-    @abstractmethod
-    def update_(self):
-        ...
+    class Config:
+        arbitrary_types_allowed = True
+        allow_mutation = True
+        extra = Extra.forbid
+
+    def __init__(self, map_fn_=None, state=list()):
+        super().__init__(
+            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,
+            )
+        else:
+            return MapMetric(
+                map_fn_=lambda *args, **kwargs: fn(self.map_fn_(*args, **kwargs)),
+                state=self.state,
+            )
+
+    def starmap(self, fn):
+        return self.map(star(fn))
+
+    def reduce(self, fn):
+        if self.map_fn_ is None:
+            return ReduceMetric(
+                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_,
+                reduce_fn=fn,
+                state=self.state,
+            )
+
+    def aggregate(self, fn):
+        return AggregateMetric(metric=self, aggregate_fn=fn)
+
+    def staraggregate(self, fn):
+        return self.aggregate(star(fn))
+
+    def update_(self, *args, **kwargs):
+        if self.map_fn_ is None:
+            self.state.append(args)
+        else:
+            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]))
+        else:
+            return self.replace(state=self.state + [self.map_fn_(*args, **kwargs)])
 
-    @abstractmethod
     def compute(self):
-        ...
+        return self.state
 
+    def log(self, tensorboard_logger, tag, step=None):
+        for name, value in self.compute().items():
+            tensorboard_logger.add_scalar(
+                f"{tag}/{name}",
+                value,
+                step,
+            )
+        return self
+
+
+Metric = MapMetric
 
-class ReduceMetric(FunctionalBase, Metric):
-    reduce_fn: Callable
-    compute_fn: Callable
-    state: Optional[Any]
+
+class ReduceMetric(BaseModel):
+    map_fn_: Callable[..., Any]
+    reduce_fn: Callable[..., Any]
+    state: Any
 
     class Config:
+        arbitrary_types_allowed = True
         allow_mutation = True
+        extra = Extra.forbid
 
-    def __init__(self, reduce_fn, compute_fn=None, initial_state=None):
-        super().__init__(
-            reduce_fn=reduce_fn,
-            compute_fn=((lambda x: x) if compute_fn is None else compute_fn),
-            state=initial_state,
-        )
+    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))
+        return self
 
     def update(self, *args, **kwargs):
-        return self.replace(state=self.reduce_fn(self.state, *args, **kwargs))
+        return self.replace(
+            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):
+        for name, value in self.compute().items():
+            tensorboard_logger.add_scalar(
+                f"{tag}/{name}",
+                value,
+                step,
+            )
+        return self
+
+
+class AggregateMetric(BaseModel):
+    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))
+        )
+
+    def starmap(self, fn):
+        return self.map(star(fn))
 
     def update_(self, *args, **kwargs):
-        self.state = self.reduce_fn(self.state, *args, **kwargs)
+        self.metric = self.metric.update(*args, **kwargs)
         return self
 
+    def update(self, *args, **kwargs):
+        return self.replace(metric=self.metric.update(*args, **kwargs))
+
     def compute(self):
-        return self.compute_fn(self.state)
+        return self.aggregate_fn(self.metric.compute())
 
-    def log(self, tensorboard_logger, tag, name, step=1):
-        tensorboard_logger.add_scalar(
-            f"{tag}/{name}",
-            self.compute(),
-            step,
-        )
+    def log(self, tensorboard_logger, tag, step=None):
+        for name, value in self.compute().items():
+            tensorboard_logger.add_scalar(
+                f"{tag}/{name}",
+                value,
+                step,
+            )
         return self
 
 
-def MapMetric(map_fn, compute_fn=np.mean):
-    """Metric version of `compute_fn(map(map_fn, input))`"""
-    return ReduceMetric(
-        reduce_fn=lambda state, *args, **kwargs: state + [map_fn(*args, **kwargs)],
-        compute_fn=compute_fn,
-        initial_state=list(),
-    )
+def test_metric():
+    pass

lantern/metric_table.py

@@ -1,12 +1,15 @@
 import textwrap
 import pandas as pd
-from lantern import FunctionalBase, Metric
-from typing import Dict
+from lantern import FunctionalBase
+from typing import Dict, Union, Any
+
+# from wire_damage.tools import MapMetric, ReduceMetric, AggregateMetric
 
 
 class MetricTable(FunctionalBase):
     name: str
-    metrics: Dict[str, Metric]
+    metrics: Dict[str, Any]
+    # metrics: Dict[str, Union[MapMetric, ReduceMetric, AggregateMetric]]
 
     class Config:
         arbitrary_types_allowed = True
@@ -18,7 +21,11 @@ def __init__(self, name, metrics):
         )
 
     def compute(self):
-        return {name: metric.compute() for name, metric in self.metrics.items()}
+        return {
+            metric_name: value
+            for metrics in self.metrics.values()
+            for metric_name, value in metrics.compute().items()
+        }
 
     def table(self):
         return "\n".join(

lantern/progress_bar.py

@@ -13,5 +13,9 @@ def ProgressBar(data_loader, name, metrics: Optional[Dict[str, Metric]] = None):
             for item in tqdm_:
                 yield item
                 tqdm_.set_postfix(
-                    {name: metric.compute() for name, metric in metrics.items()}
+                    {
+                        name: value
+                        for metrics in metrics.values()
+                        for name, value in metrics.compute().items()
+                    }
                 )

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.ReduceMetric(lambda state, loss: loss.item()),
+        loss=lantern.Metric().reduce(lambda state, loss: dict(loss=loss.item())),
     )
 
     for epoch in lantern.Epochs(2):
@@ -73,14 +73,14 @@ def test_mnist():
             train_metrics["loss"].update_(loss)
             sleep(0.5)
 
-            for name, metric in train_metrics.items():
-                metric.log(tensorboard_logger, "train", name, epoch)
+            for metrics in train_metrics.values():
+                metrics.log(tensorboard_logger, "train", epoch)
 
         print(lantern.MetricTable("train", train_metrics))
 
         evaluate_metrics = {
             name: dict(
-                loss=lantern.MapMetric(lambda loss: loss.item()),
+                loss=lantern.Metric().reduce(lambda state, loss: dict(loss=loss.item())),
             )
             for name in evaluate_data_loaders
         }
@@ -93,13 +93,13 @@ def test_mnist():
 
                 evaluate_metrics[name]["loss"].update_(loss)
 
-            for metric_name, metric in evaluate_metrics[name].items():
-                metric.log(tensorboard_logger, name, metric_name, epoch)
+            for metrics in evaluate_metrics[name].values():
+                metrics.log(tensorboard_logger, name, epoch)
 
             print(lantern.MetricTable(name, evaluate_metrics[name]))
 
         early_stopping = early_stopping.score(
-            -evaluate_metrics["evaluate_early_stopping"]["loss"].compute()
+            -evaluate_metrics["evaluate_early_stopping"]["loss"].compute()["loss"]
         )
         if early_stopping.scores_since_improvement == 0:
             torch.save(model.state_dict(), "model.pt")