merged main into solveig-branch

Author: sot176

.gitignore

@@ -1,13 +1,17 @@
 __pycache__/
 .ipynb_checkpoints/
-Data/
-Results/
-Experiments/
+Data/*
+Results/*
+Experiments/*
 _build/
-bin/
-wandb/
+bin/*
+wandb/*
 wandb_api.py
 
+#Magnus specific
+docker/*
+job*
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

environment.yml

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

main.py

@@ -1,13 +1,11 @@
-from pathlib import Path
-
 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
 
 
@@ -24,6 +22,7 @@ def main():
     ------
 
     """
+
     args = get_args()
 
     createfolders(args.datafolder, args.resultfolder, args.modelfolder)
@@ -105,18 +104,22 @@ def main():
             optimizer.step()
             optimizer.zero_grad(set_to_none=True)
 
-            preds = th.argmax(logits, dim=1)
-            metrics(y, preds)
+            metrics(y, logits)
 
             break
         print(metrics.accumulate())
         print("Dry run completed successfully.")
-        exit(0)
-
-    wandb.login(key=WANDB_API)
-    wandb.init(entity="ColabCode", project="Jan", tags=[args.modelname, args.dataset])
+        exit()
+
+    # wandb.login(key=WANDB_API)
+    wandb.init(
+            entity="ColabCode-org",
+            # entity="FYS-8805 Exam",
+            project="Test", 
+            tags=[args.modelname, args.dataset]
+            )
     wandb.watch(model)
-
+    exit()
     for epoch in range(args.epoch):
         # Training loop start
         trainingloss = []
@@ -132,8 +135,7 @@ def main():
             optimizer.zero_grad(set_to_none=True)
             trainingloss.append(loss.item())
 
-            preds = th.argmax(logits, dim=1)
-            metrics(y, preds)
+            metrics(y, logits)
 
         wandb.log(metrics.accumulate(str_prefix="Train "))
         metrics.reset()
@@ -148,8 +150,7 @@ def main():
                 loss = criterion(logits, y)
                 evalloss.append(loss.item())
 
-                preds = th.argmax(logits, dim=1)
-                metrics(y, preds)
+                metrics(y, logits)
 
         wandb.log(metrics.accumulate(str_prefix="Evaluation "))
         metrics.reset()

utils/arg_parser.py

@@ -35,7 +35,8 @@ def get_args():
 
     parser.add_argument(
         "--download-data",
-        action="store_true",
+        type=bool,
+        default=False,
         help="Whether the data should be downloaded or not. Might cause code to start a bit slowly.",
     )
 
@@ -44,14 +45,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.",
     )
 
@@ -63,6 +70,21 @@ def get_args():
         nargs="+",
         help="Which metric to use for evaluation",
     )
+    
+    parser.add_argument(
+        '--imagesize',
+        type=int,
+        default=28,
+        help='Imagesize'
+    )
+    
+    parser.add_argument(
+        '--nr_channels',
+        type=int,
+        default=1,
+        choices=[1,3],
+        help='Number of image channels'
+    )
 
     # Training specific values
     parser.add_argument(
@@ -95,4 +117,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/__init__.py

@@ -1,5 +1,6 @@
-__all__ = ["USPSDataset0_6", "USPSH5_Digit_7_9_Dataset", "MNISTDataset0_3"]
+__all__ = ["USPSDataset0_6", "USPSH5_Digit_7_9_Dataset", "MNISTDataset0_3", "SVHNDataset"]
 
 from .mnist_0_3 import MNISTDataset0_3
 from .usps_0_6 import USPSDataset0_6
 from .uspsh5_7_9 import USPSH5_Digit_7_9_Dataset
+from .svhn import SVHNDataset

utils/dataloaders/svhn.py

@@ -1,12 +1,79 @@
+import os
+import numpy as np
+from scipy.io import loadmat
 from torch.utils.data import Dataset
+from torchvision.datasets import SVHN
 
 
-class SVHN(Dataset):
-    def __init__(self):
+class SVHNDataset(Dataset):
+    def __init__(
+        self, 
+        data_path: str, 
+        train: bool,
+        transform=None, 
+        download:bool=True, 
+        nr_channels=3
+        ):
+        """
+        Initializes the SVHNDataset object.
+        Args:
+            data_path (str): Path to where the data lies. If download_data is set to True, this is where the data will be downloaded.
+            transforms: Torch composite of transformations which are to be applied to the dataset images.
+            download_data (bool): If True, downloads the dataset to the specified data_path.
+            split (str): The dataset split to use, either 'train' or 'test'.
+        Raises:
+            AssertionError: If the split is not 'train' or 'test'.
+        """
         super().__init__()
+        # assert split == "train" or split == "test"
+        self.split = 'train' if train else 'test'
+    
+        if download:
+            self._download_data(data_path)
+
+        data = loadmat(os.path.join(data_path, f"{self.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.labels[self.labels == 10] = 0
+        
+        self.nr_channels = nr_channels
+        self.transforms = transform
+
+    def _download_data(self, path: 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=self.split, download=True)
 
     def __len__(self):
-        return
+        """
+        Returns the number of samples in the dataset.
+        Returns:
+            int: The number of samples.
+        """
+        return len(self.labels)
 
     def __getitem__(self, index):
-        return
+        """
+        Retrieves the image and label at the specified index.
+        Args:
+            index (int): The index of the sample to retrieve.
+        Returns:
+            tuple: A tuple containing the image and its corresponding label.
+        """
+        img, lab = self.images[index], self.labels[index]
+
+        if self.nr_channels == 1:
+            img = np.mean(img, axis=2, keepdims=True)
+        
+        if self.transforms is not None:
+            img = self.transforms(img)
+
+        return img, lab

utils/load_data.py

@@ -1,6 +1,6 @@
 from torch.utils.data import Dataset
 
-from .dataloaders import MNISTDataset0_3, USPSDataset0_6, USPSH5_Digit_7_9_Dataset
+from .dataloaders import MNISTDataset0_3, USPSDataset0_6, USPSH5_Digit_7_9_Dataset, SVHNDataset
 
 
 def load_data(dataset: str, *args, **kwargs) -> Dataset:
@@ -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":
+            return SVHNDataset(*args, **kwargs)
+        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,9 +1,31 @@
 import torch.nn as nn
+from scipy.stats import entropy
 
 
 class EntropyPrediction(nn.Module):
-    def __init__(self):
+    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'.
+        Raises:
+            AssertionError: If the averages parameter is not 'average' or 'sum'.
+        """
         super().__init__()
 
-    def __call__(self, y_true, y_false):
-        return
+        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.
+        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.
+        """
+        entropy_values = entropy(y_true, qk=y_false_logits)
+        return entropy_values