Merge pull request #56 from SFI-Visual-Intelligence/Jan-input_checking

Added basic input checking

Author: hzavadil98

environment.yml

@@ -19,6 +19,7 @@ dependencies:
   - ruff
   - scalene
   - tqdm
+  - scipy
   - pip:
     - torch
     - torchvision

main.py

@@ -3,11 +3,11 @@
 import numpy as np
 import torch as th
 import torch.nn as nn
-import wandb
 from torch.utils.data import DataLoader
 from torchvision import transforms
 from tqdm import tqdm
 
+import wandb
 from utils import MetricWrapper, createfolders, get_args, load_data, load_model
 
 
@@ -27,7 +27,6 @@ def main():
 
     args = get_args()
 
-
     createfolders(args.datafolder, args.resultfolder, args.modelfolder)
 
     device = args.device

utils/arg_parser.py

@@ -44,14 +44,20 @@ def get_args():
         "--modelname",
         type=str,
         default="MagnusModel",
-        choices=["MagnusModel", "ChristianModel", "SolveigModel", "JanModel"],
+        choices=[
+            "MagnusModel",
+            "ChristianModel",
+            "SolveigModel",
+            "JanModel",
+            "JohanModel",
+        ],
         help="Model which to be trained on",
     )
     parser.add_argument(
         "--dataset",
         type=str,
         default="svhn",
-        choices=["svhn", "usps_0-6", "uspsh5_7_9", "mnist_0-3"],
+        choices=["svhn", "usps_0-6", "usps_7-9", "mnist_0-3", "mnist_4-9"],
         help="Which dataset to train the model on.",
     )
 
@@ -95,4 +101,10 @@ def get_args():
         action="store_true",
         help="If true, the code will not run the training loop.",
     )
-    return parser.parse_args()
+    args = parser.parse_args()
+
+    assert args.epoch > 0, "Epoch should be a positive integer."
+    assert args.learning_rate > 0, "Learning rate should be a positive float."
+    assert args.batchsize > 0, "Batch size should be a positive integer."
+
+    return args

utils/dataloaders/svhn.py

@@ -1,14 +1,14 @@
-from torch.utils.data import Dataset
+import os
+
 from scipy.io import loadmat
-import os 
+from torch.utils.data import Dataset
 from torchvision.datasets import SVHN
 
+
 class SVHNDataset(Dataset):
-    def __init__(self,
-                 datapath: str,
-                 transforms=None,
-                 download_data=True,
-                 split='train'):
+    def __init__(
+        self, datapath: str, transforms=None, download_data=True, split="train"
+    ):
         """
         Initializes the SVHNDataset object.
         Args:
@@ -20,36 +20,38 @@ def __init__(self,
             AssertionError: If the split is not 'train' or 'test'.
         """
         super().__init__()
-        assert split == 'train' or split == 'test'
-        
+        assert split == "train" or split == "test"
+
         if download_data:
             self._download_data(datapath, split)
-        
-        data = loadmat(os.path.join(datapath, f'{split}_32x32.mat'))
-        
+
+        data = loadmat(os.path.join(datapath, f"{split}_32x32.mat"))
+
         # Images on the form N x H x W x C
-        self.images = data['X'].transpose(3, 1, 0, 2)
-        self.labels = data['y'].flatten()
+        self.images = data["X"].transpose(3, 1, 0, 2)
+        self.labels = data["y"].flatten()
         self.labels[self.labels == 10] = 0
-        
+
         self.transforms = transforms
+
     def _download_data(self, path: str, split: str):
         """
         Downloads the SVHN dataset.
         Args:
             path (str): The directory where the dataset will be downloaded.
             split (str): The dataset split to download, either 'train' or 'test'.
         """
-        print(f'Downloading SVHN data into {path}')
-        SVHN(path, split=split, download=True) 
-    
+        print(f"Downloading SVHN data into {path}")
+        SVHN(path, split=split, download=True)
+
     def __len__(self):
         """
         Returns the number of samples in the dataset.
         Returns:
             int: The number of samples.
         """
         return len(self.labels)
+
     def __getitem__(self, index):
         """
         Retrieves the image and label at the specified index.
@@ -59,8 +61,8 @@ def __getitem__(self, index):
             tuple: A tuple containing the image and its corresponding label.
         """
         img, lab = self.images[index], self.labels[index]
-        
+
         if self.transforms is not None:
             img = self.transforms(img)
-        
-        return img, lab
+
+        return img, lab

utils/load_data.py

@@ -40,5 +40,9 @@ def load_data(dataset: str, *args, **kwargs) -> Dataset:
             return MNISTDataset0_3(*args, **kwargs)
         case "usps_7-9":
             return USPSH5_Digit_7_9_Dataset(*args, **kwargs)
+        case "svhn":
+            raise NotImplementedError("SVHN dataset not yet implemented.")
+        case "mnist_4-9":
+            raise NotImplementedError("MNIST 4-9 dataset not yet implemented.")
         case _:
             raise NotImplementedError(f"Dataset: {dataset} not implemented.")

utils/load_metric.py

@@ -7,7 +7,6 @@
 
 
 class MetricWrapper(nn.Module):
-
     """
     Wrapper class for metrics, that runs multiple metrics on the same data.
 
@@ -46,9 +45,7 @@ class MetricWrapper(nn.Module):
     {'entropy': [], 'f1': [], 'precision': []}
     """
 
-
     def __init__(self, *metrics, num_classes):
-
         super().__init__()
         self.metrics = {}
         self.num_classes = num_classes

utils/load_model.py

@@ -1,6 +1,6 @@
 import torch.nn as nn
 
-from .models import ChristianModel, JanModel, MagnusModel, SolveigModel
+from .models import ChristianModel, JanModel, JohanModel, MagnusModel, SolveigModel
 
 
 def load_model(modelname: str, *args, **kwargs) -> nn.Module:
@@ -44,6 +44,8 @@ def load_model(modelname: str, *args, **kwargs) -> nn.Module:
             return JanModel(*args, **kwargs)
         case "solveigmodel":
             return SolveigModel(*args, **kwargs)
+        case "johanmodel":
+            return JohanModel(*args, **kwargs)
         case _:
             errmsg = (
                 f"Model: {modelname} not implemented. "

utils/metrics/EntropyPred.py

@@ -1,24 +1,23 @@
-import numpy as np
 import torch.nn as nn
 from scipy.stats import entropy
 
 
 class EntropyPrediction(nn.Module):
-    def __init__(self, averages: str = 'average'):
+    def __init__(self, averages: str = "average"):
         """
         Initializes the EntropyPrediction module.
         Args:
-            averages (str): Specifies the method of aggregation for entropy values. 
+            averages (str): Specifies the method of aggregation for entropy values.
                             Must be either 'average' or 'sum'.
         Raises:
             AssertionError: If the averages parameter is not 'average' or 'sum'.
         """
         super().__init__()
-        
-        assert averages == 'average' or averages == 'sum'
+
+        assert averages == "average" or averages == "sum"
         self.averages = averages
         self.stored_entropy_values = []
-        
+
     def __call__(self, y_true, y_false_logits):
         """
         Computes the entropy between true labels and predicted logits, storing the results.
@@ -29,4 +28,4 @@ def __call__(self, y_true, y_false_logits):
             Appends the computed entropy values to the stored_entropy_values list.
         """
         entropy_values = entropy(y_true, qk=y_false_logits)
-        return entropy_values
+        return entropy_values

utils/models/__init__.py

@@ -1,6 +1,7 @@
-__all__ = ["MagnusModel", "ChristianModel", "JanModel", "SolveigModel"]
+__all__ = ["MagnusModel", "ChristianModel", "JanModel", "SolveigModel", "JohanModel"]
 
 from .christian_model import ChristianModel
 from .jan_model import JanModel
+from .johan_model import JohanModel
 from .magnus_model import MagnusModel
 from .solveig_model import SolveigModel

utils/models/magnus_model.py

@@ -2,43 +2,38 @@
 
 
 class MagnusModel(nn.Module):
-    def __init__(self,
-                 imagesize: int,
-                 imagechannels: int,
-                 n_classes:int=10):
-        
+    def __init__(self, imagesize: int, imagechannels: int, n_classes: int = 10):
         """
-        Magnus model contains the model for Magnus' part of the homeexam. 
+        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
         ----
             imagesize (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
-            n_classes (int): Number of classes we are to provide. 
+            n_classes (int): Number of classes we are to provide.
 
         Returns
         -------
             MagnusModel (nn.Module): Neural network as described above in this docstring.
         """
-        
         super().__init__()
-        self.imagesize = imagesize 
+        self.imagesize = imagesize
         self.imagechannels = imagechannels
-        
-        self.layer1 = nn.Sequential(*([
-            nn.Linear(self.imagechannels*self.imagesize*self.imagesize, 133),
-            nn.ReLU()
-        ]))
-        self.layer2 = nn.Sequential(*([
-            nn.Linear(133, 133),
-            nn.ReLU()
-        ]))
-        self.layer3 = nn.Sequential(*([
-            nn.Linear(133, n_classes),
-            nn.ReLU()
-        ]))
+
+        self.layer1 = nn.Sequential(
+            *(
+                [
+                    nn.Linear(
+                        self.imagechannels * self.imagesize * self.imagesize, 133
+                    ),
+                    nn.ReLU(),
+                ]
+            )
+        )
+        self.layer2 = nn.Sequential(*([nn.Linear(133, 133), nn.ReLU()]))
+        self.layer3 = nn.Sequential(*([nn.Linear(133, n_classes), nn.ReLU()]))
 
     def forward(self, x):
         """
@@ -47,17 +42,17 @@ def forward(self, x):
         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
         """
         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