Fixed some tampering with my own code.

Author: Seilmast

main.py

@@ -1,5 +1,3 @@
-from pathlib import Path
-
 import numpy as np
 import torch as th
 import torch.nn as nn
@@ -111,7 +109,7 @@ def main():
             break
         print(metrics.accumulate())
         print("Dry run completed successfully.")
-        exit(0)
+        exit()
 
     # wandb.login(key=WANDB_API)
     wandb.init(

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

@@ -7,8 +7,11 @@
 
 class SVHNDataset(Dataset):
     def __init__(
-        self, datapath: str, transforms=None, download_data=True, split="train"
-    ):
+        self, datapath: str, 
+        transforms=None, 
+        download_data=True, 
+        split="train"
+        ):
         """
         Initializes the SVHNDataset object.
         Args:

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:
@@ -41,7 +41,7 @@ def load_data(dataset: str, *args, **kwargs) -> Dataset:
         case "usps_7-9":
             return USPSH5_Digit_7_9_Dataset(*args, **kwargs)
         case "svhn":
-            raise NotImplementedError("SVHN dataset not yet implemented.")
+            raise SVHNDataset(*args, **kwargs)
         case "mnist_4-9":
             raise NotImplementedError("MNIST 4-9 dataset not yet implemented.")
         case _:

utils/models/magnus_model.py

@@ -22,18 +22,18 @@ def __init__(self, imagesize: int, imagechannels: int, n_classes: int = 10):
         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):
         """