Updated test_metrics and utility scripts

Author: sot176

tests/test_metrics.py

@@ -1,4 +1,4 @@
-from utils.metrics import Recall
+from utils.metrics import Recall, F1Score
 
 
 def test_recall():
@@ -14,3 +14,19 @@ def test_recall():
     assert recall_score.allclose(torch.tensor(0.7143), atol=1e-5), (
         f"Recall Score: {recall_score.item()}"
     )
+
+
+def test_f1score():
+    import torch
+
+    f1_metric = F1Score(num_classes=3)
+    preds = torch.tensor(
+        [[0.8, 0.1, 0.1], [0.2, 0.7, 0.1], [0.2, 0.3, 0.5], [0.1, 0.2, 0.7]]
+    )
+
+    target = torch.tensor([0, 1, 0, 2])
+
+    f1_metric.update(preds, target)
+    assert f1_metric.tp.sum().item() > 0, "Expected some true positives."
+    assert f1_metric.fp.sum().item() > 0, "Expected some false positives."
+    assert f1_metric.fn.sum().item() > 0, "Expected some false negatives."

utils/dataloaders/__init__.py

@@ -1,3 +1,4 @@
-__all__ = ["USPSDataset0_6"]
+__all__ = ["USPSDataset0_6", "USPSH5_Digit_7_9_Dataset"]
 
 from .usps_0_6 import USPSDataset0_6
+from .uspsh5_7_9 import USPSH5_Digit_7_9_Dataset

utils/load_data.py

@@ -1,11 +1,13 @@
 from torch.utils.data import Dataset
 
-from .dataloaders import USPSDataset0_6
+from .dataloaders import USPSDataset0_6, USPSH5_Digit_7_9_Dataset
 
 
 def load_data(dataset: str, *args, **kwargs) -> Dataset:
     match dataset.lower():
         case "usps_0-6":
             return USPSDataset0_6(*args, **kwargs)
+        case "usps_7-9":
+            return  USPSH5_Digit_7_9_Dataset(*args, **kwargs)
         case _:
             raise ValueError(f"Dataset: {dataset} not implemented.")

utils/load_metric.py

@@ -3,7 +3,7 @@
 import numpy as np
 import torch.nn as nn
 
-from .metrics import EntropyPrediction
+from .metrics import EntropyPrediction, F1Score
 
 
 class MetricWrapper(nn.Module):
@@ -35,7 +35,7 @@ def _get_metric(self, key):
             case "entropy":
                 return EntropyPrediction()
             case "f1":
-                raise NotImplementedError("F1 score not implemented yet")
+                raise F1Score()
             case "recall":
                 raise NotImplementedError("Recall score not implemented yet")
             case "precision":

utils/load_model.py

@@ -1,6 +1,6 @@
 import torch.nn as nn
 
-from .models import ChristianModel, MagnusModel
+from .models import ChristianModel, MagnusModel, SolveigModel
 
 
 def load_model(modelname: str, *args, **kwargs) -> nn.Module:
@@ -9,6 +9,8 @@ def load_model(modelname: str, *args, **kwargs) -> nn.Module:
             return MagnusModel(*args, **kwargs)
         case "christianmodel":
             return ChristianModel(*args, **kwargs)
+        case "solveigmodel":
+            return SolveigModel(*args, **kwargs)
         case _:
             raise ValueError(
                 f"Model: {modelname} has not been implemented. \nCheck the documentation for implemented metrics, or check your spelling"

utils/metrics/F1.py

@@ -85,16 +85,3 @@ def compute(self):
 
         return f1_score
 
-
-def test_f1score():
-    f1_metric = F1Score(num_classes=3)
-    preds = torch.tensor(
-        [[0.8, 0.1, 0.1], [0.2, 0.7, 0.1], [0.2, 0.3, 0.5], [0.1, 0.2, 0.7]]
-    )
-
-    target = torch.tensor([0, 1, 0, 2])
-
-    f1_metric.update(preds, target)
-    assert f1_metric.tp.sum().item() > 0, "Expected some true positives."
-    assert f1_metric.fp.sum().item() > 0, "Expected some false positives."
-    assert f1_metric.fn.sum().item() > 0, "Expected some false negatives."

utils/metrics/__init__.py

@@ -1,4 +1,4 @@
-__all__ = ["EntropyPrediction", "Recall", "F1"]
+__all__ = ["EntropyPrediction", "Recall", "F1Score"]
 
 from .EntropyPred import EntropyPrediction
 from .recall import Recall