Merge branch 'Jan-dataloader' into Jan-metrics

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,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
 
 
@@ -27,35 +25,25 @@ def main():
 
     args = get_args()
 
-
     createfolders(args.datafolder, args.resultfolder, args.modelfolder)
 
     device = args.device
 
     if args.dataset.lower() in ["usps_0-6", "uspsh5_7_9"]:
-        augmentations = transforms.Compose(
+        transform = transforms.Compose(
             [
                 transforms.Resize((16, 16)),
                 transforms.ToTensor(),
             ]
         )
     else:
-        augmentations = transforms.Compose([transforms.ToTensor()])
+        transform = transforms.Compose([transforms.ToTensor()])
 
-    # Dataset
-    traindata = load_data(
-        args.dataset,
-        train=True,
-        data_path=args.datafolder,
-        download=args.download_data,
-        transform=augmentations,
-    )
-    validata = load_data(
+    traindata, validata, testdata = load_data(
         args.dataset,
-        train=False,
-        data_path=args.datafolder,
-        download=args.download_data,
-        transform=augmentations,
+        data_dir=args.datafolder,
+        transform=transform,
+        val_size=args.val_size,
     )
 
     metrics = MetricWrapper(*args.metric, num_classes=traindata.num_classes, macro_averaging=args.macro_averaging)
@@ -83,6 +71,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)
@@ -107,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 = []
@@ -134,36 +129,50 @@ 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()
 
-        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)
+                metrics(y, logits)
 
-        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()

tests/test_dataloaders.py

@@ -17,18 +17,25 @@ def test_uspsdataset0_6():
 
         # Create a h5 file
         with h5py.File(tf, "w") as f:
+            targets = np.array([6, 5, 4, 3, 2, 1, 0, 0, 0, 0])
+            indices = np.arange(len(targets))
             # Populate the file with data
             f["train/data"] = np.random.rand(10, 16 * 16)
-            f["train/target"] = np.array([6, 5, 4, 3, 2, 1, 0, 0, 0, 0])
+            f["train/target"] = targets
 
         trans = transforms.Compose(
             [
-                transforms.Resize((16, 16)),  # At least for USPS
+                transforms.Resize((16, 16)),
                 transforms.ToTensor(),
             ]
         )
-        dataset = USPSDataset0_6(data_path=tempdir, train=True, transform=trans)
+        dataset = USPSDataset0_6(
+            data_path=tempdir,
+            sample_ids=indices,
+            train=True,
+            transform=trans,
+        )
         assert len(dataset) == 10
         data, target = dataset[0]
         assert data.shape == (1, 16, 16)
-        assert all(target == np.array([0, 0, 0, 0, 0, 0, 1]))
+        assert target == 6

tests/test_models.py

@@ -32,4 +32,3 @@ def test_jan_model(image_shape, num_classes):
     y = model(x)
 
     assert y.shape == (n, num_classes), f"Shape: {y.shape}"
-

utils/arg_parser.py

@@ -33,28 +33,33 @@ def get_args():
         help="Whether model should be saved or not.",
     )
 
-    parser.add_argument(
-        "--download-data",
-        action="store_true",
-        help="Whether the data should be downloaded or not. Might cause code to start a bit slowly.",
-    )
-
     # Data/Model specific values
     parser.add_argument(
         "--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(
+        "--val_size",
+        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,
@@ -70,6 +75,16 @@ def get_args():
     )
 
 
+    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(
         "--epoch",
@@ -101,4 +116,10 @@ def get_args():
         action="store_true",
         help="If the flag is included, 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,13 @@
-__all__ = ["USPSDataset0_6", "USPSH5_Digit_7_9_Dataset", "MNISTDataset0_3"]
+__all__ = [
+    "USPSDataset0_6",
+    "USPSH5_Digit_7_9_Dataset",
+    "MNISTDataset0_3",
+    "Downloader",
+    "SVHNDataset",
+]
 
+from .download import Downloader
 from .mnist_0_3 import MNISTDataset0_3
+from .svhn import SVHNDataset
 from .usps_0_6 import USPSDataset0_6
 from .uspsh5_7_9 import USPSH5_Digit_7_9_Dataset

utils/dataloaders/datasources.py

@@ -17,3 +17,26 @@
         "8ea070ee2aca1ac39742fdd1ef5ed118",
     ],
 }
+
+MNIST_SOURCE = {
+    "train_images": [
+        "https://storage.googleapis.com/cvdf-datasets/mnist/train-images-idx3-ubyte.gz",
+        "train-images-idx3-ubyte",
+        None,
+    ],
+    "train_labels": [
+        "https://storage.googleapis.com/cvdf-datasets/mnist/train-labels-idx1-ubyte.gz",
+        "train-labels-idx1-ubyte",
+        None,
+    ],
+    "test_images": [
+        "https://storage.googleapis.com/cvdf-datasets/mnist/t10k-images-idx3-ubyte.gz",
+        "t10k-images-idx3-ubyte",
+        None,
+    ],
+    "test_labels": [
+        "https://storage.googleapis.com/cvdf-datasets/mnist/t10k-labels-idx1-ubyte.gz",
+        "t10k-labels-idx1-ubyte",
+        None,
+    ],
+}