Merge branch 'main' into Jan-dataloader

Author: hzavadil98

.gitignore

@@ -1,15 +1,16 @@
 __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__/

environment.yml

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

main.py

@@ -1,11 +1,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
 
 
@@ -98,18 +98,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 = []
@@ -125,8 +129,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()
@@ -141,8 +144,7 @@ def main():
                 loss = criterion(logits, y)
                 valloss.append(loss.item())
 
-                preds = th.argmax(logits, dim=1)
-                metrics(y, preds)
+                metrics(y, logits)
 
         wandb.log(metrics.accumulate(str_prefix="Validation "))
         metrics.reset()

utils/arg_parser.py

@@ -38,14 +38,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.",
     )
     parser.add_argument(
@@ -62,6 +68,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(
@@ -94,4 +115,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

@@ -3,9 +3,11 @@
     "USPSH5_Digit_7_9_Dataset",
     "MNISTDataset0_3",
     "Downloader",
+    "SVHNDataset",
 ]
 
 from .download import Downloader
 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

@@ -6,6 +6,7 @@
     MNISTDataset0_3,
     USPSDataset0_6,
     USPSH5_Digit_7_9_Dataset,
+    SVHNDataset,
 )
 
 
@@ -59,6 +60,12 @@ def load_data(dataset: str, *args, **kwargs) -> tuple:
             dataset = MNISTDataset0_3
             train_labels, test_labels = downloader.mnist(data_dir=data_dir)
             labels = np.arange(4)
+        case "svhn":
+            dataset = SVHNDataset
+            train_labels, test_labels = downloader.svhn(data_dir=data_dir)
+            labels = np.arange(10)
+        case "mnist_4-9":
+            raise NotImplementedError("MNIST 4-9 dataset not yet implemented.")
         case _:
             raise NotImplementedError(f"Dataset: {dataset} not implemented.")
 

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,12 +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_logits):
-        return
+        assert averages == "average" or averages == "sum"
+        self.averages = averages
+        self.stored_entropy_values = []
 
-    def __reset__(self):
-        pass
+    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

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