fix preprocess in MetricMeter

Signed-off-by: Zhiyuan Chen <this@zyc.ai>

Author: ZhiyuanChen

danling/metrics/average_meter.py

@@ -50,6 +50,8 @@ class AverageMeter:
         >>> meter.update(0.7)
         >>> meter.val
         0.7
+        >>> meter.bat
+        0.7
         >>> meter.avg
         0.7
         >>> meter.update(0.9)
@@ -167,14 +169,11 @@ class AverageMeters(MetricsDict):
         'loss: 0.0000 (nan)\tauroc: 0.0000 (nan)\tr2: 0.0000 (nan)'
     """
 
-    def __init__(self, *args, default_factory: Type[AverageMeter] = AverageMeter, **kwargs) -> None:
-        for meter in args:
-            if not isinstance(meter, AverageMeter):
-                raise ValueError(f"Expected meter to be an instance of AverageMeter, but got {type(meter)}")
-        for name, meter in kwargs.items():
+    def __init__(self, default_factory: Type[AverageMeter] = AverageMeter, **meters) -> None:
+        for name, meter in meters.items():
             if not isinstance(meter, AverageMeter):
                 raise ValueError(f"Expected {name} to be an instance of AverageMeter, but got {type(meter)}")
-        super().__init__(*args, default_factory=default_factory, **kwargs)
+        super().__init__(default_factory=default_factory, **meters)
 
     @property
     def sum(self) -> FlatDict[str, float]:
@@ -288,4 +287,4 @@ def set(self, name: str, meter: AverageMeter | AverageMeters) -> None:  # pylint
             raise ValueError(
                 f"Expected meter to be an instance of AverageMeter or AverageMeters, but got {type(meter)}"
             )
-        super().set(name, meter)
+        super().set(name, meter, convert_mapping=False)

danling/metrics/metric_meter.py

@@ -70,11 +70,11 @@ class MetricMeter(AverageMeter):
     """
 
     metric: Callable
-    preprocess: Callable
+    preprocess: Optional[Callable] = None
     ignored_index: Optional[int] = None
 
     def __init__(
-        self, metric: Callable, preprocess: Callable = default_preprocess, ignored_index: int | None = None
+        self, metric: Callable, preprocess: Callable | None = default_preprocess, ignored_index: int | None = None
     ) -> None:
         self.metric = metric
         self.preprocess = preprocess
@@ -93,7 +93,8 @@ def update(  # type: ignore[override] # pylint: disable=W0237
             value: Value to be added to the average.
             n: Number of values to be added.
         """
-        input, target = self.preprocess(input, target, ignored_index=self.ignored_index)
+        if self.preprocess is not None:
+            input, target = self.preprocess(input, target, ignored_index=self.ignored_index)
         n = len(input)
         super().update(self.metric(input, target).item() * n, n=n)
 
@@ -135,27 +136,20 @@ class MetricMeters(AverageMeters):
     TypeError: ...update() missing 1 required positional argument: 'target'
     """
 
-    preprocess: Callable
+    preprocess = None
     ignored_index = None
 
     def __init__(
-        self, *args, preprocess: Callable = default_preprocess, ignored_index: int | None = None, **kwargs
+        self, preprocess: Callable | None = default_preprocess, ignored_index: int | None = None, **meters
     ) -> None:
         self.setattr("preprocess", preprocess)
         self.setattr("ignored_index", ignored_index)
-        for meter in args:
+        for name, meter in meters.items():
             if callable(meter):
-                meter = MetricMeter(meter, ignored_index=self.ignored_index)
-            if not isinstance(meter, MetricMeter):
-                raise ValueError(f"Expected meter to be an instance of MetricMeter, but got {type(meter)}")
-        for name, meter in kwargs.items():
-            if callable(meter):
-                kwargs[name] = meter = MetricMeter(meter, ignored_index=self.ignored_index)
+                meters[name] = meter = MetricMeter(meter, preprocess=None, ignored_index=self.ignored_index)
             if not isinstance(meter, MetricMeter):
                 raise ValueError(f"Expected {name} to be an instance of MetricMeter, but got {type(meter)}")
-        if ignored_index is not None:
-            self.setattr("ignored_index", ignored_index)
-        super().__init__(*args, default_factory=None, **kwargs)  # type: ignore[arg-type]
+        super().__init__(default_factory=None, **meters)  # type: ignore[arg-type]
 
     def update(  # type: ignore[override] # pylint: disable=W0221
         self,
@@ -170,13 +164,14 @@ def update(  # type: ignore[override] # pylint: disable=W0221
             target: Target values to compute the metrics.
         """
 
-        input, target = self.preprocess(input, target, ignored_index=self.ignored_index)
+        if self.preprocess is not None:
+            input, target = self.preprocess(input, target, ignored_index=self.ignored_index)
         for meter in self.values():
             meter.update(input, target)
 
     def set(self, name: str, meter: MetricMeter | Callable) -> None:  # type: ignore[override] # pylint: disable=W0237
         if callable(meter):
-            meter = MetricMeter(meter, ignored_index=self.ignored_index)
+            meter = MetricMeter(meter, preprocess=None, ignored_index=self.ignored_index)
         if not isinstance(meter, MetricMeter):
             raise ValueError(f"Expected meter to be an instance of MetricMeter, but got {type(meter)}")
         super().set(name, meter)
@@ -260,9 +255,13 @@ def set(  # pylint: disable=W0237
         name: str,
         metric: MetricMeter | MetricMeters | Callable,  # type: ignore[override]
     ) -> None:
-        if callable(metric):
+        from .metrics import Metrics
+
+        if isinstance(metric, Metrics):
+            metric = MetricMeters(preprocess=metric.preprocess, ignored_index=metric.ignored_index, **metric.metrics)
+        elif callable(metric):
             metric = MetricMeter(metric)
-        if not isinstance(metric, (MetricMeter, MetricMeters)):
+        elif not isinstance(metric, (MetricMeter, MetricMeters)):
             raise ValueError(
                 f"Expected {metric} to be an instance of MetricMeter or MetricMeters, but got {type(metric)}"
             )

danling/metrics/metrics.py

@@ -119,7 +119,7 @@ class Metrics(Metric):
     """
 
     metrics: FlatDict[str, Callable]
-    preprocess: Callable
+    preprocess: Optional[Callable] = None
     ignored_index: Optional[int] = None
     _input: Tensor
     _target: Tensor
@@ -182,6 +182,7 @@ def update(self, input: Tensor | NestedTensor | Sequence, target: Tensor | Neste
             else:
                 raise ValueError(f"Unknown input and target: {input}, {target}")
             self.flatten = True
+        # If the batch has only 1 sample at the end of an epoch, we need to flatten the input and target
         elif self.flatten:
             target = target.flatten()
             input = input.flatten() if input.numel() == target.numel() else input.view(*target.shape, -1)
@@ -196,6 +197,8 @@ def update(self, input: Tensor | NestedTensor | Sequence, target: Tensor | Neste
         if self.world_size > 1:
             input, target = self._sync(input), self._sync(target)
         input, target = input.detach().to(self.device), target.detach().to(self.device)
+        if self.preprocess is not None:
+            input, target = self.preprocess(input, target, ignored_index=self.ignored_index)
         self._input = input
         self._target = target
         self._inputs = torch.cat([self._inputs, input]).to(input.dtype)
@@ -228,9 +231,8 @@ def calculate(self, input: Tensor, target: Tensor) -> NestedDict[str, flist | fl
         ):
             return NestedDict({name: nan for name in self.metrics.keys()})
         ret = NestedDict()
-        input, target = self.preprocess(input, target, ignored_index=self.ignored_index)
         for name, metric in self.metrics.items():
-            score = self._calculate(metric, input, target, preprocess=False)
+            score = self._calculate(metric, input, target)
             if isinstance(score, Mapping):
                 if self.merge_dict:
                     ret.merge(score)
@@ -242,9 +244,7 @@ def calculate(self, input: Tensor, target: Tensor) -> NestedDict[str, flist | fl
         return ret
 
     @torch.inference_mode()
-    def _calculate(self, metric, input: Tensor, target: Tensor, preprocess: bool = True) -> flist | float:
-        if preprocess:
-            input, target = self.preprocess(input, target, ignored_index=self.ignored_index)
+    def _calculate(self, metric, input: Tensor, target: Tensor) -> flist | float:
         score = metric(input, target)
         if isinstance(score, Tensor):
             return score.item() if score.numel() == 1 else flist(score.tolist())

tests/metrics/test_metrics.py

@@ -209,7 +209,7 @@ def test_nested_tensor_regression(self):
             targets.extend(target_list)
             pred_nt, target_nt = NestedTensor(pred_list), NestedTensor(target_list)
             pred, target = torch.cat(pred_list), torch.cat(target_list)
-            metrics.update(pred_nt, target_nt)
+            metrics.update(pred_nt.tensor, target_nt)
             for metric in metric_map.values():
                 metric.update(pred.view(-1, num_outputs), target.view(-1, num_outputs))
             value, average = metrics.value(), metrics.average()
@@ -278,7 +278,7 @@ def test_nested_tensor_binary(self):
             targets.extend(target_list)
             pred_nt, target_nt = NestedTensor(pred_list), NestedTensor(target_list)
             pred, target = torch.cat(pred_list), torch.cat(target_list)
-            metrics.update(pred_nt, target_nt)
+            metrics.update(pred_nt, target_nt.tensor)
             merge_metrics.update(pred_nt, target_nt)
             for metric in metric_map.values():
                 metric.update(pred, target)