load_data - changed to accomodate train/val/test split, added test loop

Author: hzavadil98

main.py

@@ -43,16 +43,27 @@ def main():
         augmentations = transforms.Compose([transforms.ToTensor()])
 
     # Dataset
+    assert args.validation_split_percentage < 1.0 and args.validation_split_percentage > 0, "Validation split should be in interval (0,1)"
     traindata = load_data(
         args.dataset,
-        train=True,
+        split="train",
+        split_percentage=args.validation_split_percentage,
         data_path=args.datafolder,
         download=args.download_data,
         transform=augmentations,
     )
     validata = load_data(
         args.dataset,
-        train=False,
+        split="validation",
+        split_percentage=args.validation_split_percentage,
+        data_path=args.datafolder,
+        download=args.download_data,
+        transform=augmentations,
+    )
+    testdata = load_data(
+        args.dataset,
+        split="test",
+        split_percentage=args.validation_split_percentage,
         data_path=args.datafolder,
         download=args.download_data,
         transform=augmentations,
@@ -83,6 +94,9 @@ def main():
     valiloader = DataLoader(
         validata, batch_size=args.batchsize, shuffle=False, pin_memory=True
     )
+    testloader = DataLoader(
+        testdata, batch_size=args.batchsize, shuffle=False, pin_memory=True
+    )
 
     criterion = nn.CrossEntropyLoss()
     optimizer = th.optim.Adam(model.parameters(), lr=args.learning_rate)
@@ -140,30 +154,45 @@ def main():
         wandb.log(metrics.accumulate(str_prefix="Train "))
         metrics.reset()
 
-        evalloss = []
-        # Eval loop start
+        valloss = []
+        # Validation loop start
         model.eval()
         with th.no_grad():
             for x, y in tqdm(valiloader, desc="Validation"):
                 x, y = x.to(device), y.to(device)
                 logits = model.forward(x)
                 loss = criterion(logits, y)
-                evalloss.append(loss.item())
+                valloss.append(loss.item())
 
                 preds = th.argmax(logits, dim=1)
                 metrics(y, preds)
 
-        wandb.log(metrics.accumulate(str_prefix="Evaluation "))
+        wandb.log(metrics.accumulate(str_prefix="Validation "))
         metrics.reset()
 
         wandb.log(
             {
                 "Epoch": epoch,
                 "Train loss": np.mean(trainingloss),
-                "Evaluation Loss": np.mean(evalloss),
+                "Validation loss": np.mean(valloss),
             }
         )
+    
+    testloss = []
+    model.eval()
+    with th.no_grad():
+        for x, y in tqdm(testloader, desc="Testing"):
+            x, y = x.to(device), y.to(device)
+            logits = model.forward(x)
+            loss = criterion(logits, y)
+            testloss.append(loss.item())
+            
+            preds = th.argmax(logits, dim=1)
+            metrics(y, preds)
 
+    wandb.log(metrics.accumulate(str_prefix="Test "))
+    metrics.reset()
+    wandb.log({"Test loss": np.mean(testloss)})
 
 if __name__ == "__main__":
     main()

utils/arg_parser.py

@@ -54,7 +54,12 @@ def get_args():
         choices=["svhn", "usps_0-6", "uspsh5_7_9", "mnist_0-3"],
         help="Which dataset to train the model on.",
     )
-
+    parser.add_argument(
+        "--validation_split_percentage",
+        type=float,
+        default=0.2,
+        help="Percentage of training dataset to be used as validation dataset - must be within (0,1).",
+    )
     parser.add_argument(
         "--metric",
         type=str,

utils/dataloaders/mnist_0_3.py

@@ -2,9 +2,10 @@
 import os
 import urllib.request
 from pathlib import Path
+import torch
 
 import numpy as np
-from torch.utils.data import Dataset
+from torch.utils.data import Dataset, random_split
 
 
 class MNISTDataset0_3(Dataset):
@@ -59,20 +60,25 @@ class MNISTDataset0_3(Dataset):
 
     def __init__(
         self,
+        split: str,
+        split_percentage: float,
         data_path: Path,
-        train: bool = False,
-        transform=None,
         download: bool = False,
+        transform=None,
     ):
         super().__init__()
 
         self.data_path = data_path
         self.mnist_path = self.data_path / "MNIST"
-        self.train = train
+        self.split = split
+        self.split_percentage = split_percentage
         self.transform = transform
         self.download = download
         self.num_classes = 4
 
+        if self.split == "train" or self.split == "validation":
+            train = True        # used to decide whether to load training or test dataset
+        
         if not self.download and not self._chech_is_downloaded():
             raise ValueError(
                 "Data not found. Set --download-data=True to download the data."
@@ -87,13 +93,18 @@ def __init__(
             "train-labels-idx1-ubyte" if train else "t10k-labels-idx1-ubyte"
         )
 
-        labels = self._parse_labels(train=self.train)
-
+        labels = self._parse_labels()
+        
         self.idx = np.where(labels < 4)[0]
-
+        
+        if self.split != "test":
+            generator1 = torch.Generator().manual_seed(42)
+            tr, val = random_split(self.idx, [1-self.split_percentage, self.split_percentage], generator=generator1)
+            self.idx = tr if self.split == "train" else val
+                
         self.length = len(self.idx)
 
-    def _parse_labels(self, train):
+    def _parse_labels(self):
         with open(self.labels_path, "rb") as f:
             data = np.frombuffer(f.read(), dtype=np.uint8, offset=8)
         return data