Merge pull request #51 from SFI-Visual-Intelligence/main

Sync

Author: Seilmast

.github/workflows/test.yml

@@ -0,0 +1,29 @@
+name: Test
+
+on:
+  push:
+    branches: [ main ]
+  pull_request:
+    branches: [ main ]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+
+    steps:
+    - name: Checkout
+      uses: actions/checkout@v4
+
+    - uses: mamba-org/setup-micromamba@v1
+      with:
+        micromamba-version: '2.0.5-0' # any version from https://github.com/mamba-org/micromamba-releases
+        environment-file: environment.yml
+        init-shell: bash
+        cache-environment: true
+        post-cleanup: 'all'
+        generate-run-shell: false
+
+    - name: Run tests
+      run: |
+        PYTHONPATH=. pytest tests
+      shell: bash -el {0}

.gitignore

@@ -5,6 +5,8 @@ Results/
 Experiments/
 _build/
 bin/
+wandb/
+wandb_api.py
 
 #Magnus specific
 docker/*

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,5 +18,9 @@ dependencies:
   - pytest
   - ruff
   - scalene
+  - tqdm
+  - pip:
+    - torch
+    - torchvision
 prefix: /opt/miniconda3/envs/cc-exam
 

main.py

@@ -1,13 +1,14 @@
-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
+from utils import MetricWrapper, createfolders, get_args, load_data, load_model
 
 
 def main():
@@ -23,202 +24,146 @@ 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"],
-        help="Model which to be trained on",
-    )
-    parser.add_argument(
-        "--dataset",
-        type=str,
-        default="svhn",
-        choices=["svhn", "usps_0-6"],
-        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",
-    )
-    parser.add_argument('--imagesize',
-                        type=int,
-                        default=28,
-                        help='Size of images')
-    parser.add_argument('--imagechannels',
-                        type=int,
-                        default=1,
-                        choices=[1,3],
-                        help='Number of color channels in the image.')
-    
-    
-    
-
-    # 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 = get_args()
 
-    args = parser.parse_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,
     )
 
-    # Find number of channels in the dataset
-    if len(traindata[0][0].shape) == 2:
-        channels = 1
-    else:
-        channels = traindata[0][0].shape[0]
+    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:
+        data_shape = (1, *data_shape)
 
     # load model
     model = load_model(
         args.modelname,
-        in_channels=channels,
+        image_shape=data_shape,
         num_classes=traindata.num_classes,
     )
     model.to(device)
 
-    trainloader = DataLoader(traindata,
-                             batch_size=args.batchsize,
-                             shuffle=True,
-                             pin_memory=True,
-                             drop_last=True)
-    valiloader = DataLoader(validata,
-                            batch_size=args.batchsize,
-                            shuffle=False,
-                            pin_memory=True)
+    trainloader = DataLoader(
+        traindata,
+        batch_size=args.batchsize,
+        shuffle=True,
+        pin_memory=True,
+        drop_last=True,
+    )
+    valiloader = DataLoader(
+        validata, batch_size=args.batchsize, shuffle=False, pin_memory=True
+    )
 
     criterion = nn.CrossEntropyLoss()
     optimizer = th.optim.Adam(model.parameters(), lr=args.learning_rate)
 
     # This allows us to load all the components without running the training loop
     if args.dry_run:
-        print("Dry run completed")
+        dry_run_loader = DataLoader(
+            traindata,
+            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)
+            logits = model.forward(x)
+
+            loss = criterion(logits, y)
+            loss.backward()
+
+            optimizer.step()
+            optimizer.zero_grad(set_to_none=True)
+
+            preds = th.argmax(logits, dim=1)
+            metrics(y, preds)
+
+            break
+        print(metrics.accumulate())
+        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):
-
         # 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())
 
+            preds = th.argmax(logits, dim=1)
+            metrics(y, preds)
+
+        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())
 
-        wandb.log({
-            'Epoch': epoch,
-            'Train loss': np.mean(trainingloss),
-            'Evaluation Loss': np.mean(evalloss)
-        })
+                preds = th.argmax(logits, dim=1)
+                metrics(y, preds)
+
+        wandb.log(metrics.accumulate(str_prefix="Evaluation "))
+        metrics.reset()
+
+        wandb.log(
+            {
+                "Epoch": epoch,
+                "Train loss": np.mean(trainingloss),
+                "Evaluation Loss": np.mean(evalloss),
+            }
+        )
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()

pyproject.toml

@@ -0,0 +1,3 @@
+[tool.isort]
+profile = "black"
+line_length = 88