Merge branch 'main' into dataset-rework

c-salomonsen · c-salomonsen · commit 11bbfd99133e · 2025-02-04T13:27:01.000+01:00
diff --git a/main.py b/main.py
@@ -66,14 +66,14 @@ def main():
         "--modelname",
         type=str,
         default="MagnusModel",
-        choices=["MagnusModel", "ChristianModel"],
+        choices=["MagnusModel", "ChristianModel", "SolveigModel"],
         help="Model which to be trained on",
     )
     parser.add_argument(
         "--dataset",
         type=str,
         default="svhn",
-        choices=["svhn", "usps_0-6"],
+        choices=["svhn", "usps_0-6", "uspsh5_7_9"],
         help="Which dataset to train the model on.",
     )
 
diff --git a/tests/test_metrics.py b/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."
diff --git a/utils/dataloaders/__init__.py b/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
diff --git a/utils/load_data.py b/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.")
diff --git a/utils/load_metric.py b/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, precision
 
 
 class MetricWrapper(nn.Module):
@@ -35,11 +35,11 @@ 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":
-                raise NotImplementedError("Precision score not implemented yet")
+                return precision()
             case "accuracy":
                 raise NotImplementedError("Accuracy score not implemented yet")
             case _:
diff --git a/utils/load_model.py b/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"
diff --git a/utils/metrics/F1.py b/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."
diff --git a/utils/metrics/__init__.py b/utils/metrics/__init__.py
@@ -1,4 +1,5 @@
-__all__ = ["EntropyPrediction", "Recall"]
+__all__ = ["EntropyPrediction", "Recall", "F1Score"]
 
 from .EntropyPred import EntropyPrediction
+from .F1 import F1Score
 from .recall import Recall
diff --git a/utils/metrics/precision.py b/utils/metrics/precision.py
@@ -0,0 +1,105 @@
+import torch
+import torch.nn as nn
+
+USE_MEAN = True
+
+# Precision = TP / (TP + FP)
+
+
+class Precision(nn.Module):
+    """Metric module for precision. Can calculate precision both as a mean of precisions or as brute function of true positives and false positives. This is for now controller with the USE_MEAN macro.
+
+    Parameters
+    ----------
+    num_classes : int
+        Number of classes in the dataset.
+    """
+
+    def __init__(self, num_classes):
+        super().__init__()
+
+        self.num_classes = num_classes
+
+    def forward(self, y_true, y_pred):
+        """Calculates the precision score given number of classes and the true and predicted labels.
+
+        Parameters
+        ----------
+        y_true : torch.tensor
+            true labels
+        y_pred : torch.tensor
+            predicted labels
+
+        Returns
+        -------
+        torch.tensor
+            precision score
+        """
+        # One-hot encode the target tensor
+        true_oh = torch.zeros(y_true.size(0), self.num_classes).scatter_(
+            1, y_true.unsqueeze(1), 1
+        )
+        pred_oh = torch.zeros(y_pred.size(0), self.num_classes).scatter_(
+            1, y_pred.unsqueeze(1), 1
+        )
+
+        if USE_MEAN:
+            tp = torch.sum(true_oh * pred_oh, 0)
+            fp = torch.sum(~true_oh.bool() * pred_oh, 0)
+
+        else:
+            tp = torch.sum(true_oh * pred_oh)
+            fp = torch.sum(~true_oh[pred_oh.bool()].bool())
+
+        return torch.nanmean(tp / (tp + fp))
+
+
+def test_precision_case1():
+    true_precision = 25.0 / 36 if USE_MEAN else 7.0 / 10
+
+    true1 = torch.tensor([0, 1, 2, 1, 0, 2, 1, 0, 2, 1])
+    pred1 = torch.tensor([0, 2, 1, 1, 0, 2, 0, 0, 2, 1])
+    P = Precision(3)
+    precision1 = P(true1, pred1)
+    assert precision1.allclose(torch.tensor(true_precision), atol=1e-5), (
+        f"Precision Score: {precision1.item()}"
+    )
+
+
+def test_precision_case2():
+    true_precision = 8.0 / 15 if USE_MEAN else 6.0 / 15
+
+    true2 = torch.tensor([0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4])
+    pred2 = torch.tensor([0, 0, 4, 3, 4, 0, 4, 4, 2, 3, 4, 1, 2, 4, 0])
+    P = Precision(5)
+    precision2 = P(true2, pred2)
+    assert precision2.allclose(torch.tensor(true_precision), atol=1e-5), (
+        f"Precision Score: {precision2.item()}"
+    )
+
+
+def test_precision_case3():
+    true_precision = 3.0 / 4 if USE_MEAN else 4.0 / 5
+
+    true3 = torch.tensor([0, 0, 0, 1, 0])
+    pred3 = torch.tensor([1, 0, 0, 1, 0])
+    P = Precision(2)
+    precision3 = P(true3, pred3)
+    assert precision3.allclose(torch.tensor(true_precision), atol=1e-5), (
+        f"Precision Score: {precision3.item()}"
+    )
+
+
+def test_for_zero_denominator():
+    true_precision = 0.0
+    true4 = torch.tensor([1, 1, 1, 1, 1])
+    pred4 = torch.tensor([0, 0, 0, 0, 0])
+    P = Precision(2)
+    precision4 = P(true4, pred4)
+    assert precision4.allclose(torch.tensor(true_precision), atol=1e-5), (
+        f"Precision Score: {precision4.item()}"
+    )
+
+
+if __name__ == "__main__":
+    pass
diff --git a/utils/models/__init__.py b/utils/models/__init__.py
@@ -1,4 +1,5 @@
-__all__ = ["MagnusModel", "ChristianModel"]
+__all__ = ["MagnusModel", "ChristianModel", "SolveigModel"]
 
 from .christian_model import ChristianModel
 from .magnus_model import MagnusModel
+from .solveig_model import SolveigModel
diff --git a/utils/models/johan_model.py b/utils/models/johan_model.py
@@ -0,0 +1,62 @@
+import torch.nn as nn
+
+"""
+Multi-layer perceptron model for image classification.
+"""
+
+# class NeuronLayer(nn.Module):
+#     def __init__(self, in_features, out_features):
+#         super().__init__()
+
+#         self.fc = nn.Linear(in_features, out_features)
+#         self.relu = nn.ReLU()
+
+#     def forward(self, x):
+#         x = self.fc(x)
+#         x = self.relu(x)
+#         return x
+
+
+class JohanModel(nn.Module):
+    """Small MLP model for image classification.
+
+    Parameters
+    ----------
+    in_features : int
+        Numer of input features.
+    num_classes : int
+        Number of classes in the dataset.
+
+    """
+
+    def __init__(self, image_shape, num_classes):
+        super().__init__()
+
+        # Extract features from image shape
+        self.in_channels = image_shape[0]
+        self.height = image_shape[1]
+        self.width = image_shape[2]
+        self.num_classes = num_classes
+        self.in_features = self.in_channels * self.height * self.width
+
+        self.fc1 = nn.Linear(self.in_features, 77)
+        self.fc2 = nn.Linear(77, 77)
+        self.fc3 = nn.Linear(77, 77)
+        self.fc4 = nn.Linear(77, num_classes)
+        self.softmax = nn.Softmax(dim=1)
+        self.relu = nn.ReLU()
+
+    def forward(self, x):
+        for layer in [self.fc1, self.fc2, self.fc3, self.fc4]:
+            x = layer(x)
+            x = self.relu(x)
+        x = self.softmax(x)
+        return x
+
+
+# TODO
+# Add your tests here
+
+
+if __name__ == "__main__":
+    pass  # Add your tests here
diff --git a/utils/models/solveig_model.py b/utils/models/solveig_model.py
@@ -0,0 +1,74 @@
+import torch
+import torch.nn as nn
+
+
+class SolveigModel(nn.Module):
+    """
+        A Convolutional Neural Network model for classification.
+
+         Args
+        ----
+        image_shape : tuple(int, int, int)
+            Shape of the input image (C, H, W).
+        num_classes : int
+            Number of classes in the dataset.
+
+        Attributes:
+        -----------
+        conv_block1 : nn.Sequential
+            First convolutional block containing a convolutional layer, ReLU activation, and max-pooling.
+        conv_block2 : nn.Sequential
+            Second convolutional block containing a convolutional layer and ReLU activation.
+        conv_block3 : nn.Sequential
+            Third convolutional block containing a convolutional layer and ReLU activation.
+        fc1 : nn.Linear
+            Fully connected layer that outputs the final classification scores.
+        """
+
+    def __init__(self, image_shape, num_classes):
+        super().__init__()
+
+        C, *_ = image_shape
+
+        # Define the first convolutional block (conv + relu + maxpool)
+        self.conv_block1 = nn.Sequential(
+            nn.Conv2d(in_channels=C, out_channels=25, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.MaxPool2d(kernel_size=2, stride=2)
+        )
+
+        # Define the second convolutional block (conv + relu)
+        self.conv_block2 = nn.Sequential(
+            nn.Conv2d(in_channels=25, out_channels=50, kernel_size=3, padding=1),
+            nn.ReLU()
+        )
+
+        # Define the third convolutional block (conv + relu)
+        self.conv_block3 = nn.Sequential(
+            nn.Conv2d(in_channels=50, out_channels=100, kernel_size=3, padding=1),
+            nn.ReLU()
+        )
+
+        self.fc1 = nn.Linear(100 * 8 * 8, num_classes)
+
+    def forward(self, x):
+        x = self.conv_block1(x)
+        x = self.conv_block2(x)
+        x = self.conv_block3(x)
+        x = torch.flatten(x, 1)
+
+        x = self.fc1(x)
+        x = nn.Softmax(x)
+
+        return x
+
+
+if __name__ == "__main__":
+
+    x = torch.randn(1,3, 16, 16)
+
+    model = SolveigModel(x.shape[1:], 3)
+
+    y = model(x)
+
+    print(y)
