Merge branch 'main' into christian/sphinx-autoapi

Author: c-salomonsen

.gitignore

@@ -5,6 +5,8 @@ Results/
 Experiments/
 _build/
 bin/
+wandb/
+wandb_api.py
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

main.py

@@ -1,4 +1,3 @@
-import argparse
 from pathlib import Path
 
 import numpy as np
@@ -9,7 +8,7 @@
 from torchvision import transforms
 from tqdm import tqdm
 
-from utils import MetricWrapper, createfolders, load_data, load_model
+from utils import MetricWrapper, createfolders, get_args, load_data, load_model
 
 
 def main():
@@ -25,113 +24,21 @@ 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",
-        action="store_true",
-        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"],
-        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)
-
-    augmentations = transforms.Compose(
-        [
-            transforms.Resize((16, 16)),  # At least for USPS
-            transforms.ToTensor(),
-        ]
-    )
+    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(
@@ -149,6 +56,8 @@ def main():
         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:
@@ -180,28 +89,32 @@ def main():
     if args.dry_run:
         dry_run_loader = DataLoader(
             traindata,
-            batch_size=1,
+            batch_size=20,
             shuffle=True,
             pin_memory=True,
             drop_last=True,
         )
 
         for x, y in tqdm(dry_run_loader, desc="Dry run", total=1):
             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)
 
-            break
+            preds = th.argmax(logits, dim=1)
+            metrics(y, preds)
 
+            break
+        print(metrics.__getmetrics__())
         print("Dry run completed successfully.")
         exit(0)
 
-    wandb.init(project="", tags=[])
+    wandb.login(key=WANDB_API)
+    wandb.init(entity="ColabCode", project="Jan", tags=[args.modelname, args.dataset])
     wandb.watch(model)
 
     for epoch in range(args.epoch):
@@ -210,25 +123,37 @@ def main():
         model.train()
         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())
 
+            preds = th.argmax(logits, dim=1)
+            metrics(y, preds)
+
+        wandb.log(metrics.__getmetrics__(str_prefix="Train "))
+        metrics.__resetvalues__()
+
         evalloss = []
         # Eval loop start
         model.eval()
         with th.no_grad():
             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())
 
+                preds = th.argmax(logits, dim=1)
+                metrics(y, preds)
+
+        wandb.log(metrics.__getmetrics__(str_prefix="Evaluation "))
+        metrics.__resetvalues__()
+
         wandb.log(
             {
                 "Epoch": epoch,

tests/test_metrics.py

@@ -1,4 +1,4 @@
-from utils.metrics import F1Score, Recall
+from utils.metrics import Accuracy, F1Score, Precision, Recall
 
 
 def test_recall():
@@ -30,3 +30,70 @@ def test_f1score():
     assert f1_metric.tp.sum().item() > 0, "Expected some true positives."
     assert f1_metric.fp.sum().item() > 0, "Expected some false positives."
     assert f1_metric.fn.sum().item() > 0, "Expected some false negatives."
+
+
+def test_precision_case1():
+    import torch
+
+    for boolean, true_precision in zip([True, False], [25.0 / 36, 7.0 / 10]):
+        true1 = torch.tensor([0, 1, 2, 1, 0, 2, 1, 0, 2, 1])
+        pred1 = torch.tensor([0, 2, 1, 1, 0, 2, 0, 0, 2, 1])
+        P = Precision(3, use_mean=boolean)
+        precision1 = P(true1, pred1)
+        assert precision1.allclose(torch.tensor(true_precision), atol=1e-5), (
+            f"Precision Score: {precision1.item()}"
+        )
+
+
+def test_precision_case2():
+    import torch
+
+    for boolean, true_precision in zip([True, False], [8.0 / 15, 6.0 / 15]):
+        true2 = torch.tensor([0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 1, 2, 3, 4])
+        pred2 = torch.tensor([0, 0, 4, 3, 4, 0, 4, 4, 2, 3, 4, 1, 2, 4, 0])
+        P = Precision(5, use_mean=boolean)
+        precision2 = P(true2, pred2)
+        assert precision2.allclose(torch.tensor(true_precision), atol=1e-5), (
+            f"Precision Score: {precision2.item()}"
+        )
+
+
+def test_precision_case3():
+    import torch
+
+    for boolean, true_precision in zip([True, False], [3.0 / 4, 4.0 / 5]):
+        true3 = torch.tensor([0, 0, 0, 1, 0])
+        pred3 = torch.tensor([1, 0, 0, 1, 0])
+        P = Precision(2, use_mean=boolean)
+        precision3 = P(true3, pred3)
+        assert precision3.allclose(torch.tensor(true_precision), atol=1e-5), (
+            f"Precision Score: {precision3.item()}"
+        )
+
+
+def test_for_zero_denominator():
+    import torch
+
+    for boolean in [True, False]:
+        true4 = torch.tensor([1, 1, 1, 1, 1])
+        pred4 = torch.tensor([0, 0, 0, 0, 0])
+        P = Precision(2, use_mean=boolean)
+        precision4 = P(true4, pred4)
+        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