added pyproject and mnist dataloader, although not finished yet

Author: Johanmkr

.gitignore

@@ -1,10 +1,16 @@
 __pycache__/
 .ipynb_checkpoints/
-Data/
-Results/
-Experiments/
+Data/*
+Results/*
+Experiments/*
 _build/
-bin/
+bin/*
+wandb/*
+wandb_api.py
+
+#Magnus specific
+docker/*
+job*
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

doc/about.md

@@ -1,3 +1,3 @@
 # About this code
 
-Work in progress ...
+Work is still in progress ...

doc/conf.py

@@ -7,8 +7,12 @@
 
 extensions = [
     "myst_parser",  # in order to use markdown
+    "autoapi.extension",  # in order to generate API documentation
 ]
 
+# search this directory for Python files
+autoapi_dirs = ["../utils"]
+
 myst_enable_extensions = [
     "colon_fence",  # ::: can be used instead of ``` for better rendering
 ]

environment.yml

@@ -18,6 +18,8 @@ dependencies:
   - pytest
   - ruff
   - scalene
+  - tqdm
+  - scipy
   - pip:
     - torch
     - torchvision

main.py

@@ -1,13 +1,12 @@
-import argparse
-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
 
-from utils import MetricWrapper, createfolders, load_data, load_model
+import wandb
+from utils import MetricWrapper, createfolders, get_args, load_data, load_model
 
 
 def main():
@@ -23,122 +22,41 @@ def main():
     ------
 
     """
-    parser = argparse.ArgumentParser(
-        prog="",
-        description="",
-        epilog="",
-    )
-    # Structuture related values
-    parser.add_argument(
-        "--datafolder",
-        type=Path,
-        default="Data",
-        help="Path to where data will be saved during training.",
-    )
-    parser.add_argument(
-        "--resultfolder",
-        type=Path,
-        default="Results",
-        help="Path to where results will be saved during evaluation.",
-    )
-    parser.add_argument(
-        "--modelfolder",
-        type=Path,
-        default="Experiments",
-        help="Path to where model weights will be saved at the end of training.",
-    )
-    parser.add_argument(
-        "--savemodel",
-        type=bool,
-        default=False,
-        help="Whether model should be saved or not.",
-    )
-
-    parser.add_argument(
-        "--download-data",
-        type=bool,
-        default=False,
-        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"],
-        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"],
-        help="Which dataset to train the model on.",
-    )
-
-    parser.add_argument(
-        "--metric",
-        type=str,
-        default=["entropy"],
-        choices=["entropy", "f1", "recall", "precision", "accuracy"],
-        nargs="+",
-        help="Which metric to use for evaluation",
-    )
-
-    # Training specific values
-    parser.add_argument(
-        "--epoch",
-        type=int,
-        default=20,
-        help="Amount of training epochs the model will do.",
-    )
-    parser.add_argument(
-        "--learning_rate",
-        type=float,
-        default=0.001,
-        help="Learning rate parameter for model training.",
-    )
-    parser.add_argument(
-        "--batchsize",
-        type=int,
-        default=64,
-        help="Amount of training images loaded in one go",
-    )
-    parser.add_argument(
-        "--device",
-        type=str,
-        default="cpu",
-        choices=["cuda", "cpu", "mps"],
-        help="Which device to run the training on.",
-    )
-    parser.add_argument(
-        "--dry_run",
-        action="store_true",
-        help="If true, the code will not run the training loop.",
-    )
-
-    args = parser.parse_args()
+    args = get_args()
 
     createfolders(args.datafolder, args.resultfolder, args.modelfolder)
 
     device = args.device
 
-    metrics = MetricWrapper(*args.metric)
+    if args.dataset.lower() in ["usps_0-6", "uspsh5_7_9"]:
+        augmentations = transforms.Compose(
+            [
+                transforms.Resize((16, 16)),
+                transforms.ToTensor(),
+            ]
+        )
+    else:
+        augmentations = 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(
         args.dataset,
         train=False,
         data_path=args.datafolder,
+        download=args.download_data,
+        transform=augmentations,
     )
 
+    metrics = MetricWrapper(*args.metric, num_classes=traindata.num_classes)
+
     # Find the shape of the data, if is 2D, add a channel dimension
     data_shape = traindata[0][0].shape
     if len(data_shape) == 2:
@@ -168,37 +86,75 @@ def main():
 
     # This allows us to load all the components without running the training loop
     if args.dry_run:
-        print("Dry run completed")
-        exit(0)
+        dry_run_loader = DataLoader(
+            traindata,
+            batch_size=20,
+            shuffle=True,
+            pin_memory=True,
+            drop_last=True,
+        )
 
-    wandb.init(project="", tags=[])
-    wandb.watch(model)
+        for x, y in tqdm(dry_run_loader, desc="Dry run", total=1):
+            x, y = x.to(device), y.to(device)
+            logits = model.forward(x)
 
+            loss = criterion(logits, y)
+            loss.backward()
+
+            optimizer.step()
+            optimizer.zero_grad(set_to_none=True)
+
+            metrics(y, logits)
+
+            break
+        print(metrics.accumulate())
+        print("Dry run completed successfully.")
+        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 = []
         model.train()
-        for x, y in trainloader:
+        for x, y in tqdm(trainloader, desc="Training"):
             x, y = x.to(device), y.to(device)
-            pred = model.forward(x)
+            logits = model.forward(x)
 
-            loss = criterion(y, pred)
+            loss = criterion(logits, y)
             loss.backward()
 
             optimizer.step()
             optimizer.zero_grad(set_to_none=True)
             trainingloss.append(loss.item())
 
+            metrics(y, logits)
+
+        wandb.log(metrics.accumulate(str_prefix="Train "))
+        metrics.reset()
+
         evalloss = []
         # Eval loop start
         model.eval()
         with th.no_grad():
-            for x, y in valiloader:
+            for x, y in tqdm(valiloader, desc="Validation"):
                 x, y = x.to(device), y.to(device)
-                pred = model.forward(x)
-                loss = criterion(y, pred)
+                logits = model.forward(x)
+                loss = criterion(logits, y)
                 evalloss.append(loss.item())
 
+                metrics(y, logits)
+
+        wandb.log(metrics.accumulate(str_prefix="Evaluation "))
+        metrics.reset()
+
         wandb.log(
             {
                 "Epoch": epoch,

pyproject.toml

@@ -22,3 +22,6 @@ dependencies = [
     "torch>=2.6.0",
     "torchvision>=0.21.0",
 ]
+[tool.isort]
+profile = "black"
+line_length = 88

tests/test_dataloaders.py

@@ -3,17 +3,31 @@
 
 def test_uspsdataset0_6():
     from pathlib import Path
-    from tempfile import TemporaryFile
+    from tempfile import TemporaryDirectory
 
     import h5py
     import numpy as np
+    from torchvision import transforms
 
-    with TemporaryFile() as tf:
+    # Create a temporary directory (deleted after the test)
+    with TemporaryDirectory() as tempdir:
+        tempdir = Path(tempdir)
+
+        tf = tempdir / "usps.h5"
+
+        # Create a h5 file
         with h5py.File(tf, "w") as f:
+            # 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])
 
-        dataset = USPSDataset0_6(data_path=tf, train=True)
+        trans = transforms.Compose(
+            [
+                transforms.Resize((16, 16)),  # At least for USPS
+                transforms.ToTensor(),
+            ]
+        )
+        dataset = USPSDataset0_6(data_path=tempdir, train=True, transform=trans)
         assert len(dataset) == 10
         data, target = dataset[0]
         assert data.shape == (1, 16, 16)

tests/test_metrics.py

@@ -1,6 +1,4 @@
-
-from utils.metrics import F1Score, Precision, Recall
-
+from utils.metrics import Accuracy, F1Score, Precision, Recall
 
 
 def test_recall():
@@ -84,3 +82,18 @@ def test_for_zero_denominator():
         assert precision4.allclose(torch.tensor(0.0), atol=1e-5), (
             f"Precision Score: {precision4.item()}"
         )
+
+
+def test_accuracy():
+    import torch
+
+    accuracy = Accuracy(num_classes=5)
+
+    y_true = torch.tensor([0, 3, 2, 3, 4])
+    y_pred = torch.tensor([0, 1, 2, 3, 4])
+
+    accuracy_score = accuracy(y_true, y_pred)
+
+    assert torch.abs(torch.tensor(accuracy_score - 0.8)) < 1e-5, (
+        f"Accuracy Score: {accuracy_score.item()}"
+    )

tests/test_models.py

@@ -1,7 +1,7 @@
 import pytest
 import torch
 
-from utils.models import ChristianModel
+from utils.models import ChristianModel, JanModel
 
 
 @pytest.mark.parametrize(
@@ -17,6 +17,19 @@ def test_christian_model(image_shape, num_classes):
     y = model(x)
 
     assert y.shape == (n, num_classes), f"Shape: {y.shape}"
-    assert y.sum(dim=1).allclose(torch.ones(n), atol=1e-5), (
-        f"Softmax output should sum to 1, but got: {y.sum()}"
-    )
+
+
+@pytest.mark.parametrize(
+    "image_shape, num_classes",
+    [((1, 28, 28), 4), ((3, 16, 16), 10)],
+)
+def test_jan_model(image_shape, num_classes):
+    n, c, h, w = 5, *image_shape
+
+    model = JanModel(image_shape, num_classes)
+
+    x = torch.randn(n, c, h, w)
+    y = model(x)
+
+    assert y.shape == (n, num_classes), f"Shape: {y.shape}"
+

utils/__init__.py

@@ -1,5 +1,6 @@
-__all__ = ["createfolders", "load_data", "load_model", "MetricWrapper"]
+__all__ = ["createfolders", "load_data", "load_model", "MetricWrapper", "get_args"]
 
+from .arg_parser import get_args
 from .createfolders import createfolders
 from .load_data import load_data
 from .load_metric import MetricWrapper