feature!: refactor metrics

BREAKING CHANGE

Author: samedii

docs/source/lantern.rst

@@ -13,7 +13,7 @@ lantern package
    :members:
    :undoc-members:
    :member-order: bysource
-.. autoclass:: lantern.Metrics
+.. autoclass:: lantern.MetricTable
    :members:
    :undoc-members:
    :member-order: bysource
@@ -29,6 +29,10 @@ lantern package
    :members:
    :undoc-members:
    :member-order: bysource
+.. autoclass:: lantern.Tensor
+   :members:
+   :undoc-members:
+   :member-order: bysource
 
 .. autofunction:: lantern.Epochs
 .. autofunction:: lantern.module_device
@@ -40,4 +44,5 @@ lantern package
 .. autofunction:: lantern.update_cpu_model
 .. autofunction:: lantern.step
 .. autofunction:: lantern.to_device
-.. autofunction:: lantern.worker_init
+.. autofunction:: lantern.worker_init_fn
+.. autofunction:: lantern.set_learning_rate

lantern/__init__.py

@@ -6,8 +6,8 @@
 from lantern.functional_base import FunctionalBase
 from lantern.tensor import Tensor
 from lantern.epochs import Epochs
-from lantern.metric import ReduceMetric, MapMetric
-from lantern.metrics import Metrics
+from lantern.metric import Metric, ReduceMetric, MapMetric
+from lantern.metric_table import MetricTable
 from lantern.module_device import module_device
 from lantern.module_compose import ModuleCompose
 from lantern.to_device import to_device
@@ -16,7 +16,7 @@
 from lantern.requires_grad import requires_grad, requires_nograd
 from lantern.set_learning_rate import set_learning_rate
 from lantern.set_seeds import set_seeds
-from lantern.worker_init import worker_init
+from lantern.worker_init_fn import worker_init_fn
 from lantern.evaluate import evaluate
 from lantern.step import step
 from lantern.train import train

lantern/early_stopping.py

@@ -13,7 +13,6 @@ class EarlyStopping(FunctionalBase):
 
     class Config:
         arbitrary_types_allowed = True
-        allow_mutation = False
 
     def score(self, value):
         if self.best_score is None or value >= self.best_score:

lantern/metric.py

@@ -1,27 +1,59 @@
 import numpy as np
+from abc import ABC, abstractmethod
+from lantern import FunctionalBase
+from typing import Callable, Any, Optional
 
 
-class ReduceMetric:
+class Metric(ABC):
+    @abstractmethod
+    def update(self):
+        ...
+
+    @abstractmethod
+    def update_(self):
+        ...
+
+    @abstractmethod
+    def compute(self):
+        ...
+
+
+class ReduceMetric(FunctionalBase, Metric):
+    reduce_fn: Callable
+    compute_fn: Callable
+    state: Optional[Any]
+
+    class Config:
+        allow_mutation = True
+
     def __init__(self, reduce_fn, compute_fn=None, initial_state=None):
-        self.reduce_fn = reduce_fn
-        if compute_fn is None:
-            self.compute_fn = lambda x: x
-        else:
-            self.compute_fn = compute_fn
-        self.state = initial_state
-
-    def reduce(self, *args, **kwargs):
-        return ReduceMetric(
-            reduce_fn=self.reduce_fn,
-            compute_fn=self.compute_fn,
-            initial_state=self.reduce_fn(self.state, *args, **kwargs),
+        super().__init__(
+            reduce_fn=reduce_fn,
+            compute_fn=((lambda x: x) if compute_fn is None else compute_fn),
+            state=initial_state,
         )
 
+    def update(self, *args, **kwargs):
+        return self.replace(state=self.reduce_fn(self.state, *args, **kwargs))
+
+    def update_(self, *args, **kwargs):
+        self.state = self.reduce_fn(self.state, *args, **kwargs)
+        return self
+
     def compute(self):
         return self.compute_fn(self.state)
 
+    def log(self, tensorboard_logger, tag, name, step=1):
+        tensorboard_logger.add_scalar(
+            f"{tag}/{name}",
+            self.compute(),
+            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,

lantern/metric_table.py

@@ -0,0 +1,39 @@
+import textwrap
+import pandas as pd
+from lantern import FunctionalBase, Metric
+from typing import Dict
+
+
+class MetricTable(FunctionalBase):
+    name: str
+    metrics: Dict[str, Metric]
+
+    class Config:
+        arbitrary_types_allowed = True
+
+    def __init__(self, name, metrics):
+        super().__init__(
+            name=name,
+            metrics=metrics,
+        )
+
+    def compute(self):
+        return {name: metric.compute() for name, metric in self.metrics.items()}
+
+    def table(self):
+        return "\n".join(
+            [
+                f"{self.name}:",
+                textwrap.indent(
+                    (
+                        pd.Series(self.compute()).to_string(
+                            name=True, dtype=False, index=True
+                        )
+                    ),
+                    prefix="  ",
+                ),
+            ]
+        )
+
+    def __str__(self):
+        return self.table()

lantern/metrics.py

@@ -1,12 +1,17 @@
 from tqdm import tqdm
+from typing import Dict, Optional
+from lantern.metric import Metric
 
 
-def ProgressBar(data_loader, metrics=None):
+def ProgressBar(data_loader, name, metrics: Optional[Dict[str, Metric]] = None):
     """Simple progress bar with metrics"""
     if metrics is None:
-        return tqdm(data_loader, leave=False)
+        for item in tqdm(data_loader, desc=name, leave=False):
+            yield item
     else:
-        with tqdm(data_loader, desc=metrics.name, leave=False) as tqdm_:
+        with tqdm(data_loader, desc=name, leave=False) as tqdm_:
             for item in tqdm_:
                 yield item
-                tqdm_.set_postfix(metrics.compute())
+                tqdm_.set_postfix(
+                    {name: metric.compute() for name, metric in metrics.items()}
+                )

lantern/progress_bar.py

@@ -0,0 +1,8 @@
+from lantern import set_seeds
+
+
+def worker_init_fn(seed):
+    def worker_init(worker_id):
+        set_seeds(seed * 2 ** 16 + worker_id * 2 ** 24)
+
+    return worker_init

lantern/worker_init.py

@@ -14,6 +14,7 @@
 
 
 def test_mnist():
+    torch.set_grad_enabled(False)
 
     device = torch.device("cpu")
     model = ModuleCompose(
@@ -31,74 +32,71 @@ def test_mnist():
         [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
     )
 
-    gradient_dataset = datastream.Dataset.from_subscriptable(
+    train_dataset = datastream.Dataset.from_subscriptable(
         datasets.MNIST("data", train=True, transform=transform, download=True)
     )
     early_stopping_dataset = datastream.Dataset.from_subscriptable(
         datasets.MNIST("data", train=False, transform=transform)
     )
 
-    gradient_data_loader = (
-        datastream.Datastream(gradient_dataset).take(16 * 4).data_loader(batch_size=16)
+    train_data_loader = (
+        datastream.Datastream(train_dataset).take(16 * 4).data_loader(batch_size=16)
     )
     early_stopping_data_loader = (
         datastream.Datastream(early_stopping_dataset)
         .take(16 * 4)
         .data_loader(batch_size=16)
     )
     evaluate_data_loaders = dict(
-        evaluate_gradient=gradient_data_loader,
+        evaluate_train=train_data_loader,
         evaluate_early_stopping=early_stopping_data_loader,
     )
 
     tensorboard_logger = torch.utils.tensorboard.SummaryWriter()
     early_stopping = lantern.EarlyStopping(tensorboard_logger=tensorboard_logger)
-    gradient_metrics = lantern.Metrics(
-        name="gradient",
-        tensorboard_logger=tensorboard_logger,
-        metrics=dict(
-            loss=lantern.ReduceMetric(lambda state, examples, predictions, loss: loss),
-        ),
+    train_metrics = dict(
+        loss=lantern.ReduceMetric(lambda state, loss: loss.item()),
     )
 
     for epoch in lantern.Epochs(2):
 
-        with lantern.module_train(model):
-            for examples, targets in lantern.ProgressBar(
-                gradient_data_loader, metrics=gradient_metrics[["loss"]]
-            ):
+        for examples, targets in lantern.ProgressBar(
+            train_data_loader, "train", train_metrics
+        ):
+            with lantern.module_train(model), torch.enable_grad():
                 predictions = model(examples)
                 loss = F.nll_loss(predictions, targets)
                 loss.backward()
-                optimizer.step()
-                optimizer.zero_grad()
+            optimizer.step()
+            optimizer.zero_grad()
 
-                (
-                    gradient_metrics.update_(
-                        examples, predictions.detach(), loss.detach()
-                    ).log_()
-                )
-                sleep(0.5)
-        gradient_metrics.print()
+            train_metrics["loss"].update_(loss)
+            sleep(0.5)
+
+            for name, metric in train_metrics.items():
+                metric.log(tensorboard_logger, "train", name, epoch)
+
+        print(lantern.MetricTable("train", train_metrics))
 
         evaluate_metrics = {
-            name: lantern.Metrics(
-                name=name,
-                tensorboard_logger=tensorboard_logger,
-                metrics=dict(
-                    loss=lantern.MapMetric(lambda examples, predictions, loss: loss),
-                ),
+            name: dict(
+                loss=lantern.MapMetric(lambda loss: loss.item()),
             )
-            for name in evaluate_data_loaders.keys()
+            for name in evaluate_data_loaders
         }
 
-        with lantern.module_eval(model), torch.no_grad():
-            for name, data_loader in evaluate_data_loaders.items():
-                for examples, targets in tqdm(data_loader, desc=name, leave=False):
+        for name, data_loader in evaluate_data_loaders.items():
+            for examples, targets in tqdm(data_loader, desc=name, leave=False):
+                with lantern.module_eval(model):
                     predictions = model(examples)
                     loss = F.nll_loss(predictions, targets)
-                    evaluate_metrics[name].update_(examples, predictions, loss)
-                evaluate_metrics[name].log_().print()
+
+                evaluate_metrics[name]["loss"].update_(loss)
+
+            for metric_name, metric in evaluate_metrics[name].items():
+                metric.log(tensorboard_logger, name, metric_name, epoch)
+
+            print(lantern.MetricTable(name, evaluate_metrics[name]))
 
         early_stopping = early_stopping.score(
             -evaluate_metrics["evaluate_early_stopping"]["loss"].compute()