Added several tests, undid some changes to the argparser file

Author: Seilmast

main.py

@@ -22,13 +22,13 @@ def main():
     ------
 
     """
-
+    
     args = get_args()
-
+    
     createfolders(args.datafolder, args.resultfolder, args.modelfolder)
-
+    
     device = args.device
-
+    
     if args.dataset.lower() in ["usps_0-6", "uspsh5_7_9"]:
         augmentations = transforms.Compose(
             [
@@ -38,7 +38,7 @@ def main():
         )
     else:
         augmentations = transforms.Compose([transforms.ToTensor()])
-
+    
     # Dataset
     traindata = load_data(
         args.dataset,
@@ -54,22 +54,22 @@ def main():
         download=args.download_data,
         transform=augmentations,
     )
-
-    metrics = MetricWrapper(*args.metric, num_classes=traindata.num_classes)
-
+    
+    metrics = MetricWrapper(traindata.num_classes, *args.metric)
+    
     # Find the shape of the data, if is 2D, add a channel dimension
     data_shape = traindata[0][0].shape
     if len(data_shape) == 2:
         data_shape = (1, *data_shape)
-
+    
     # load model
     model = load_model(
         args.modelname,
         image_shape=data_shape,
         num_classes=traindata.num_classes,
     )
     model.to(device)
-
+    
     trainloader = DataLoader(
         traindata,
         batch_size=args.batchsize,

tests/test_dataloaders.py

@@ -1,5 +1,4 @@
-from utils.dataloaders.usps_0_6 import USPSDataset0_6
-
+from utils.dataloaders import USPSDataset0_6, SVHNDataset
 
 def test_uspsdataset0_6():
     from pathlib import Path
@@ -32,3 +31,33 @@ def test_uspsdataset0_6():
         data, target = dataset[0]
         assert data.shape == (1, 16, 16)
         assert all(target == np.array([0, 0, 0, 0, 0, 0, 1]))
+
+        
+def test_svhn_dataset():
+    import os
+    from tempfile import TemporaryDirectory
+    from torchvision import transforms
+    
+    with TemporaryDirectory() as tempdir: 
+        
+        trans = transforms.Compose([
+            transforms.Resize((28,28)),
+            transforms.ToTensor()
+        ])
+        
+        dataset = SVHNDataset(tempdir,
+                              train=True,
+                              transform=trans,
+                              download=True,
+                              nr_channels=1)
+        
+        assert dataset.__len__() != 0
+        assert os.path.exists(os.path.join(tempdir, 'train_32x32.mat'))
+        
+        img, label = dataset.__getitem__(0)
+        assert len(img.size()) == 3 and img.size() == (1,28,28) and img.size(0) == 1
+        assert len(label.size()) == 1
+
+
+if __name__ == '__main__':
+    test_svhn_dataset()

tests/test_metrics.py

@@ -1,4 +1,4 @@
-from utils.metrics import Accuracy, F1Score, Precision, Recall
+from utils.metrics import Accuracy, F1Score, Precision, Recall, EntropyPrediction
 
 
 def test_recall():
@@ -97,3 +97,21 @@ def test_accuracy():
     assert torch.abs(torch.tensor(accuracy_score - 0.8)) < 1e-5, (
         f"Accuracy Score: {accuracy_score.item()}"
     )
+
+def test_entropypred():
+    import torch as th 
+    
+    metric = EntropyPrediction(averages='mean')
+
+    true_lab = th.Tensor([0,1,1,2,4,3]).reshape(6,1).type(th.LongTensor)
+    pred_logits = th.nn.functional.one_hot(true_lab, 5)
+    
+    #Test for log(0) errors and expected output
+    assert th.abs((th.sum(metric(true_lab, pred_logits)) - 0.0)) < 1e-5
+    
+    pred_logits = th.rand(6,5)
+    metric2 = EntropyPrediction(averages='sum')
+    
+    #Test for averaging metric consistency
+    assert th.abs(th.sum(6*metric(true_lab, pred_logits) - metric2(true_lab, pred_logits))) < 1e-5
+    

tests/test_models.py

@@ -1,7 +1,7 @@
 import pytest
 import torch
 
-from utils.models import ChristianModel, JanModel
+from utils.models import ChristianModel, JanModel, MagnusModel
 
 
 @pytest.mark.parametrize(
@@ -33,3 +33,20 @@ def test_jan_model(image_shape, num_classes):
 
     assert y.shape == (n, num_classes), f"Shape: {y.shape}"
 
+
+@pytest.mark.parameterize(
+    "image_shape",
+    [(3,28,28)]
+)
+def test_magnus_model(image_shape):
+    import torch as th 
+    
+    n, c, h, w = 5, *image_shape
+    model = MagnusModel([h,w], 10, c)
+    
+    x = th.rand((n, c, h, w))
+    with th.no_grad():
+        y = model(x)
+        
+    assert y.shape == (n, 10), f"Shape: {y.shape}"
+    

utils/arg_parser.py

@@ -35,8 +35,7 @@ def get_args():
 
     parser.add_argument(
         "--download-data",
-        type=bool,
-        default=False,
+        action="store_true",
         help="Whether the data should be downloaded or not. Might cause code to start a bit slowly.",
     )
 

utils/dataloaders/svhn.py

@@ -40,6 +40,7 @@ def __init__(
 
         self.nr_channels = nr_channels
         self.transforms = transform
+        self.num_classes = len(np.unique(self.labels))
 
     def _download_data(self, path: str):
         """

utils/load_metric.py

@@ -1,8 +1,6 @@
 import copy
-
 import numpy as np
 import torch.nn as nn
-
 from .metrics import Accuracy, EntropyPrediction, F1Score, Precision, Recall
 
 
@@ -45,7 +43,7 @@ class MetricWrapper(nn.Module):
     {'entropy': [], 'f1': [], 'precision': []}
     """
 
-    def __init__(self, *metrics, num_classes):
+    def __init__(self, num_classes, *metrics):
         super().__init__()
         self.metrics = {}
         self.num_classes = num_classes
@@ -72,7 +70,8 @@ def _get_metric(self, key):
 
         match key.lower():
             case "entropy":
-                return EntropyPrediction(num_classes=self.num_classes)
+                #Not dependent on knowing the number of classes
+                return EntropyPrediction()
             case "f1":
                 return F1Score(num_classes=self.num_classes)
             case "recall":

utils/metrics/EntropyPred.py

@@ -8,24 +8,54 @@ def __init__(self, averages: str = "average"):
         Initializes the EntropyPrediction module.
         Args:
             averages (str): Specifies the method of aggregation for entropy values.
-                            Must be either 'average' or 'sum'.
+                            Must be either 'mean', 'sum' or 'none.
         Raises:
-            AssertionError: If the averages parameter is not 'average' or 'sum'.
+            AssertionError: If the averages parameter is not 'mean' or 'sum'.
         """
         super().__init__()
 
-        assert averages == "average" or averages == "sum"
+        assert averages == "mean" or averages == "sum"
         self.averages = averages
         self.stored_entropy_values = []
 
-    def __call__(self, y_true, y_false_logits):
+    def __call__(self, y_true, y_logits):
         """
-        Computes the entropy between true labels and predicted logits, storing the results.
+        Computes the Shannon Entropy of the predicted logits, storing the results.
         Args:
-            y_true: The true labels.
-            y_false_logits: The predicted logits.
-        Side Effects:
-            Appends the computed entropy values to the stored_entropy_values list.
+            y_true: The true labels. Does nothing, but needed for compatability sake. 
+            y_logits: The predicted logits.
         """
-        entropy_values = entropy(y_true, qk=y_false_logits)
+        entropy_values = entropy(y_logits, axis=1)
+        entropy_values = th.from_numpy(entropy_values)
+        
+        #Fix numerical errors for perfect guesses
+        entropy_values[entropy_values == th.inf] = 0
+        entropy_values = th.nan_to_num(entropy_values)
+
+        
+        if self.averages == 'mean':
+            entropy_values = th.mean(entropy_values)
+            
+        elif self.averages == 'sum':
+            entropy_values = th.sum(entropy_values) 
+        
+        elif self.averages == 'none':
+            return entropy_values
+        
+        
         return entropy_values
+
+
+if __name__ == '__main__':
+    import torch as th 
+    
+    metric = EntropyPrediction(averages='mean')
+
+    true_lab = th.Tensor([0,1,1,2,4,3]).reshape(6,1)
+    pred_logits = th.nn.functional.one_hot(true_lab, 5)
+    
+    assert th.abs((th.sum(metric(true_lab, pred_logits)) - 0.0)) < 1e-5
+    
+    pred_logits = th.rand(6,5)
+    metric2 = EntropyPrediction(averages='sum')
+    assert th.abs(th.sum(6*metric(true_lab, pred_logits) - metric2(true_lab, pred_logits))) < 1e-5

utils/models/magnus_model.py

@@ -2,28 +2,25 @@
 
 
 class MagnusModel(nn.Module):
-    def __init__(self, image_shape: int, num_classes: int, imagechannels: int):
+    def __init__(self, image_shape, num_classes: int, nr_channels: int):
         """
-        Magnus model contains the model for Magnus' part of the homeexam.
-        This class contains a neural network consisting of three linear layers of 133 neurons each,
-        with ReLU activation between each layer.
-
-        Args
-        ----
-            image_shape (int): Expected size of input image. This is needed to scale first layer input
-            imagechannels (int): Expected number of image channels. This is needed to scale first layer input
-            num_classes (int): Number of classes we are to provide.
-
-        Returns
-        -------
-            MagnusModel (nn.Module): Neural network as described above in this docstring.
+        Initializes the MagnusModel, a neural network designed for image classification tasks.
+        
+        The model consists of three linear layers, each with 133 neurons, and uses ReLU activation 
+        functions between the layers. The first layer's input size is determined by the image shape 
+        and number of channels, while the output layer's size is determined by the number of classes.
+        Args:
+            image_shape (tuple): A tuple representing the dimensions of the input image (Channels, Height, Width).
+            num_classes (int): The number of output classes for classification.
+            nr_channels (int): The number of channels in the input image.
+        Returns:
+            MagnusModel (nn.Module): An instance of the MagnusModel neural network.
         """
         super().__init__()
-        self.image_shape = image_shape
-        self.imagechannels = imagechannels
-
+        _, H, W = image_shape
+        
         self.layer1 = nn.Sequential(*([
-                        nn.Linear(self.imagechannels * self.imagesize * self.imagesize, 133),
+                        nn.Linear(nr_channels * H * W, 133),
                         nn.ReLU(),
                     ]))
         self.layer2 = nn.Sequential(*([
@@ -32,27 +29,32 @@ def __init__(self, image_shape: int, num_classes: int, imagechannels: int):
                     ]))
         self.layer3 = nn.Sequential(*([
                         nn.Linear(133, num_classes), 
-                        nn.ReLU()
                     ]))
-
     def forward(self, x):
         """
-        Forward pass of MagnusModel
-
-        Args
-        ----
-            x (th.Tensor): Four-dimensional tensor in the form (Batch Size x Channels x Image Height x Image Width)
-
-        Returns
-        -------
-            out (th.Tensor): Class-logits of network given input x
+        Defines the forward pass of the MagnusModel.
+        Args:
+            x (torch.Tensor): A four-dimensional tensor with shape (Batch Size, Channels, Image Height, Image Width).
+        Returns:
+            torch.Tensor: The output tensor containing class logits for each input sample.
         """
-        assert len(x.size) == 4
-
+        assert len(x.size()) == 4
         x = x.view(x.size(0), -1)
-
         x = self.layer1(x)
         x = self.layer2(x)
         out = self.layer3(x)
-
         return out
+
+
+if __name__ == '__main__':
+    import torch as th 
+    
+    data_shape = [28,28]
+    
+    data_shape = (3, *data_shape)
+    model = MagnusModel(data_shape, 10)
+    
+    dummy_img = th.rand((5,*data_shape))
+    print(dummy_img.size())
+    with th.no_grad():
+        print(model(dummy_img).size())