Merge pull request #40 from SFI-Visual-Intelligence/Jan-dev

Main.py updates

Author: hzavadil98

.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

@@ -87,13 +87,13 @@ def test_for_zero_denominator():
 def test_accuracy():
     import torch
 
-    accuracy = Accuracy()
+    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), (
+    assert torch.abs(torch.tensor(accuracy_score - 0.8)) < 1e-5, (
         f"Accuracy Score: {accuracy_score.item()}"
     )

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

utils/arg_parser.py

@@ -0,0 +1,98 @@
+import argparse
+from pathlib import Path
+
+
+def get_args():
+    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", "JanModel"],
+        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.",
+    )
+    return parser.parse_args()

utils/dataloaders/mnist_0_3.py

@@ -134,11 +134,11 @@ def __len__(self):
 
     def __getitem__(self, index):
         with open(self.labels_path, "rb") as f:
-            f.seek(8 + index)  # Jump to the label position
+            f.seek(8 + self.idx[index])  # Jump to the label position
             label = int.from_bytes(f.read(1), byteorder="big")  # Read 1 byte for label
 
         with open(self.images_path, "rb") as f:
-            f.seek(16 + index * 28 * 28)  # Jump to image position
+            f.seek(16 + self.idx[index] * 28 * 28)  # Jump to image position
             image = np.frombuffer(f.read(28 * 28), dtype=np.uint8).reshape(
                 28, 28
             )  # Read image data