Proposing change to metricwrapper class

Author: Seilmast

utils/load_metric.py

@@ -8,71 +8,63 @@
 
 class MetricWrapper(nn.Module):
     """
-    Wrapper class for metrics, that runs multiple metrics on the same data.
-
+    A wrapper class for evaluating multiple metrics on the same dataset.
+    This class allows you to compute several metrics simultaneously on given
+    true and predicted labels. It supports a variety of common metrics and
+    provides methods to accumulate results and reset the state.
     Args
     ----
+        num_classes : int
+            The number of classes in the classification task.
         metrics : list[str]
-            List of metrics to run on the data.
-
+            A list of metric names to be evaluated.
     Attributes
     ----------
         metrics : dict
-            Dictionary containing the metric functions.
-        tmp_scores : dict
-            Dictionary containing the temporary scores of the metrics.
-
+            A dictionary mapping metric names to their corresponding functions.
+        num_classes : int
+            The number of classes for the classification task.
     Methods
     -------
         __call__(y_true, y_pred)
-            Call the metric functions on the true and predicted labels.
-        accumulate()
-            Get the average scores of the metrics.
+            Computes the specified metrics on the provided true and predicted labels.
+        __getmetrics__(str_prefix: str = None)
+            Retrieves the computed metrics, optionally prefixed with a string.
         reset()
-            Reset the temporary scores of the metrics.
-
+            Resets the state of all metric computations.
     Examples
     --------
-    >>> from utils import MetricWrapper
-    >>> metrics = MetricWrapper("entropy", "f1", "precision")
+    >>> from utils import MetricWrapperProposed
+    >>> metrics = MetricWrapperProposed(2, "entropy", "f1", "precision")
     >>> y_true = [0, 1, 0, 1]
     >>> y_pred = [0, 1, 1, 0]
     >>> metrics(y_true, y_pred)
-    >>> metrics.accumulate()
+    >>> metrics.__getmetrics__()
     {'entropy': 0.6931471805599453, 'f1': 0.5, 'precision': 0.5}
     >>> metrics.reset()
-    >>> metrics.accumulate()
+    >>> metrics.__getmetrics__()
     {'entropy': [], 'f1': [], 'precision': []}
     """
 
     def __init__(self, num_classes, *metrics):
         super().__init__()
         self.metrics = {}
         self.num_classes = num_classes
-
         for metric in metrics:
             self.metrics[metric] = self._get_metric(metric)
 
-        self.tmp_scores = copy.deepcopy(self.metrics)
-        for key in self.tmp_scores:
-            self.tmp_scores[key] = []
-
     def _get_metric(self, key):
         """
-        Get the metric function based on the key
-
+        Retrieves the metric function based on the provided key.
         Args
         ----
-            key (str): metric name
-
+            key (str): The name of the metric.
         Returns
         -------
-            metric (callable): metric function
+            metric (callable): The function that computes the metric.
         """
-
         match key.lower():
             case "entropy":
-                # Not dependent on knowing the number of classes
                 return EntropyPrediction()
             case "f1":
                 return F1Score(num_classes=self.num_classes)
@@ -87,18 +79,17 @@ def _get_metric(self, key):
 
     def __call__(self, y_true, y_pred):
         for key in self.metrics:
-            self.tmp_scores[key].append(self.metrics[key](y_true, y_pred))
+            self.metrics[key](y_true, y_pred)
 
-    def accumulate(self, str_prefix: str = None):
+    def __getmetrics__(self, str_prefix: str = None):
         return_metrics = {}
         for key in self.metrics:
             if str_prefix is not None:
-                return_metrics[str_prefix + key] = np.mean(self.tmp_scores[key])
+                return_metrics[str_prefix + key] = self.metrics[key].__returnmetric__()
             else:
-                return_metrics[key] = np.mean(self.tmp_scores[key])
-
+                return_metrics[key] = self.metrics[key].__returnmetric__()
         return return_metrics
 
     def reset(self):
-        for key in self.tmp_scores:
-            self.tmp_scores[key] = []
+        for key in self.metrics:
+            self.metrics[key].reset()

utils/metrics/EntropyPred.py

@@ -46,8 +46,21 @@ def __call__(self, y_true, y_logits):
         elif self.averages == "none":
             return entropy_values
 
+        self.stored_entropy_values.append(entropy_values)
+
         return entropy_values
 
+    def return_value(self):
+        if self.averages == "mean":
+            self.stored_entropy_values = th.mean(self.stored_entropy_values)
+        elif self.averages == "sum":
+            self.stored_entropy_values = th.sum(self.stored_entropy_values)
+        elif self.averages == "none":
+            return self.stored_entropy_values
+
+    def reset(self):
+        self.stored_entropy_values = []
+
 
 if __name__ == "__main__":
     import torch as th