load_data now gives arguments to the datasets

Author: c-salomonsen

main.py

@@ -73,7 +73,7 @@ def main():
         "--dataset",
         type=str,
         default="svhn",
-        choices=["svhn"],
+        choices=["svhn", "usps_0-6"],
         help="Which dataset to train the model on.",
     )
 
@@ -119,8 +119,17 @@ def main():
     metrics = MetricWrapper(*args.metric)
 
     # Dataset
-    traindata = load_data(args.dataset)
-    validata = load_data(args.dataset)
+    traindata = load_data(
+        args.dataset,
+        train=True,
+        data_path=args.datafolder,
+        download=args.download_data,
+    )
+    validata = load_data(
+        args.dataset,
+        train=False,
+        data_path=args.datafolder,
+    )
 
     trainloader = DataLoader(traindata,
                              batch_size=args.batchsize,
@@ -144,7 +153,7 @@ def main():
         # Training loop start
         trainingloss = []
         model.train()
-        for x, y in traindata:
+        for x, y in trainloader:
             x, y = x.to(device), y.to(device)
             pred = model.forward(x)
 

utils/dataloaders/usps_0_6.py

@@ -17,19 +17,21 @@ class USPSDataset0_6(Dataset):
 
     Args
     ----
-    path : pathlib.Path
+    data_path : pathlib.Path
         Path to the USPS dataset file.
-    mode : str
-        Mode of the dataset. Must be either 'train' or 'test'.
+    train : bool, optional
+        Mode of the dataset.
     transform : callable, optional
         A function/transform that takes in a sample and returns a transformed version.
+    download : bool, optional
+        Whether to download the Dataset.
 
     Attributes
     ----------
     path : pathlib.Path
         Path to the USPS dataset file.
     mode : str
-        Mode of the dataset.
+        Mode of the dataset, either train or test.
     transform : callable
         A function/transform that takes in a sample and returns a transformed version.
     idx : numpy.ndarray
@@ -59,15 +61,21 @@ class USPSDataset0_6(Dataset):
     6
     """
 
-    def __init__(self, path: Path, mode: str = "train", transform=None):
+    def __init__(
+        self,
+        data_path: Path,
+        train: bool = False,
+        transform=None,
+        download: bool = False,
+    ):
         super().__init__()
-        self.path = path
-        self.mode = mode
+        self.path = list(data_path.glob("*.h5"))[0]
         self.transform = transform
 
-        if self.mode not in ["train", "test"]:
-            raise ValueError("Invalid mode. Must be either 'train' or 'test'")
+        if download:
+            raise NotImplementedError("Download functionality not implemented.")
 
+        self.mode = "train" if train else "test"
         self.idx = self._index()
 
     def _index(self):

utils/load_data.py

@@ -1,10 +1,11 @@
-from dataloaders import USPS_0_6
 from torch.utils.data import Dataset
 
+from .dataloaders import USPSDataset0_6
 
-def load_data(dataset: str) -> Dataset:
+
+def load_data(dataset: str, *args, **kwargs) -> Dataset:
     match dataset.lower():
         case "usps_0-6":
-            return USPS_0_6
+            return USPSDataset0_6(*args, **kwargs)
         case _:
             raise ValueError(f"Dataset: {dataset} not implemented.")