Added micro/macro averaging option to MetricsWrapper and as commandline argument

Author: hzavadil98

main.py

@@ -58,7 +58,7 @@ def main():
         transform=augmentations,
     )
 
-    metrics = MetricWrapper(*args.metric, num_classes=traindata.num_classes)
+    metrics = MetricWrapper(*args.metric, num_classes=traindata.num_classes, macro_averaging=args.macro_averaging)
 
     # Find the shape of the data, if is 2D, add a channel dimension
     data_shape = traindata[0][0].shape

utils/arg_parser.py

@@ -63,6 +63,12 @@ def get_args():
         nargs="+",
         help="Which metric to use for evaluation",
     )
+    parser.add_argument(
+        "--macro_averaging",
+        action="store_true",
+        help="If the flag is included, the metrics will be calculated using macro averaging.",
+    )
+        
 
     # Training specific values
     parser.add_argument(
@@ -93,6 +99,6 @@ def get_args():
     parser.add_argument(
         "--dry_run",
         action="store_true",
-        help="If true, the code will not run the training loop.",
+        help="If the flag is included, the code will not run the training loop.",
     )
     return parser.parse_args()

utils/load_metric.py

@@ -47,11 +47,12 @@ class MetricWrapper(nn.Module):
     """
 
 
-    def __init__(self, *metrics, num_classes):
+    def __init__(self, *metrics, num_classes, macro_averaging=False):
 
         super().__init__()
         self.metrics = {}
         self.num_classes = num_classes
+        self.macro_averaging = macro_averaging
 
         for metric in metrics:
             self.metrics[metric] = self._get_metric(metric)
@@ -77,13 +78,13 @@ def _get_metric(self, key):
             case "entropy":
                 return EntropyPrediction(num_classes=self.num_classes)
             case "f1":
-                return F1Score(num_classes=self.num_classes)
+                return F1Score(num_classes=self.num_classes, macro_averaging=self.macro_averaging)
             case "recall":
-                return Recall(num_classes=self.num_classes)
+                return Recall(num_classes=self.num_classes, macro_averaging=self.macro_averaging)
             case "precision":
-                return Precision(num_classes=self.num_classes)
+                return Precision(num_classes=self.num_classes, macro_averaging=self.macro_averaging)
             case "accuracy":
-                return Accuracy(num_classes=self.num_classes)
+                return Accuracy(num_classes=self.num_classes, macro_averaging=self.macro_averaging)
             case _:
                 raise ValueError(f"Metric {key} not supported")
 

utils/metrics/accuracy.py

@@ -3,10 +3,11 @@
 
 
 class Accuracy(nn.Module):
-    def __init__(self, num_classes):
+    def __init__(self, num_classes, macro_averaging=False):
         super().__init__()
         self.num_classes = num_classes
-
+        self.macro_averaging = macro_averaging
+        
     def forward(self, y_true, y_pred):
         """
         Compute the accuracy of the model.
@@ -23,12 +24,71 @@ def forward(self, y_true, y_pred):
         float
             Accuracy score.
         """
+        if y_pred.dim() > 1:
+            y_pred = y_pred.argmax(dim=1)
+        if self.macro_averaging:
+            return self._macro_acc(y_true, y_pred)
+        else:
+            return self._micro_acc(y_true, y_pred)
+    
+    def _macro_acc(self, y_true, y_pred):
+        """
+        Compute the macro-average accuracy.
+
+        Parameters
+        ----------
+        y_true : torch.Tensor
+            True labels.
+        y_pred : torch.Tensor
+            Predicted labels.
+
+        Returns
+        -------
+        float
+            Macro-average accuracy score.
+        """
+        y_true, y_pred = y_true.flatten(), y_pred.flatten()  # Ensure 1D shape
+
+        classes = torch.unique(y_true)  # Find unique class labels
+        acc_per_class = []
+        
+        for c in classes:
+            mask = (y_true == c)  # Mask for class c
+            acc = (y_pred[mask] == y_true[mask]).float().mean()  # Accuracy for class c
+            acc_per_class.append(acc)
+        
+        macro_acc = torch.stack(acc_per_class).mean().item()  # Average across classes
+        return macro_acc
+    
+    def _micro_acc(self, y_true, y_pred):
+        """
+        Compute the micro-average accuracy.
+
+        Parameters
+        ----------
+        y_true : torch.Tensor
+            True labels.
+        y_pred : torch.Tensor
+            Predicted labels.
+
+        Returns
+        -------
+        float
+            Micro-average accuracy score.
+        """
         return (y_true == y_pred).float().mean().item()
 
 
 if __name__ == "__main__":
+    accuracy = Accuracy(5)
+    macro_accuracy = Accuracy(5, macro_averaging=True)
+    
     y_true = torch.tensor([0, 3, 2, 3, 4])
     y_pred = torch.tensor([0, 1, 2, 3, 4])
-
-    accuracy = Accuracy()
     print(accuracy(y_true, y_pred))
+    print(macro_accuracy(y_true, y_pred))
+    
+    y_true = torch.tensor([0, 3, 2, 3, 4])
+    y_onehot_pred = torch.tensor([[1, 0, 0, 0, 0], [0, 1, 0, 0, 0], [0, 0, 1, 0, 0], [0, 0, 0, 1, 0], [0, 0, 0, 0, 1]])
+    print(accuracy(y_true, y_onehot_pred))
+    print(macro_accuracy(y_true, y_onehot_pred))