Merge pull request #21 from SFI-Visual-Intelligence/christian/dataloader

- Implemented USPS dataloader for integers 0-6.
- Fixed a few bugs related to relative imports (changing `from metrics import ...`to `from .metrics import ...`)
- Changed from using `os.path` to the superior `pathlib`.
- Added `--dry_run` option to cli interface.
- Added `mps` backend to `--device`.

Author: c-salomonsen

environment.yml

@@ -9,6 +9,14 @@ dependencies:
   - sphinx-autobuild
   - sphinx-rtd-theme
   - pip
+  - h5py
+  - black
+  - isort
+  - jupyterlab
+  - numpy
+  - pandas
   - pytest
+  - ruff
+  - scalene
 prefix: /opt/miniconda3/envs/cc-exam
 

main.py

@@ -1,68 +1,148 @@
-import torch as th
-import torch.nn as nn 
-from torch.utils.data import DataLoader
 import argparse
-import wandb
+from pathlib import Path
+
 import numpy as np
-from utils import MetricWrapper, load_model, load_data, createfolders
+import torch as th
+import torch.nn as nn
+import wandb
+from torch.utils.data import DataLoader
+
+from utils import MetricWrapper, createfolders, load_data, load_model
 
 
 def main():
-    '''
-    
+    """
+
     Parameters
     ----------
-    
+
     Returns
     -------
-    
+
     Raises
     ------
-    
-    '''
+
+    """
     parser = argparse.ArgumentParser(
-        prog='',
-        description='',
-        epilog='',
-    )
-    #Structuture related values
-    parser.add_argument('--datafolder', type=str, default='Data/', help='Path to where data will be saved during training.')
-    parser.add_argument('--resultfolder', type=str, default='Results/', help='Path to where results will be saved during evaluation.')
-    parser.add_argument('--modelfolder', type=str, 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'], help="Model which to be trained on")
-    parser.add_argument('--dataset', type=str, default='svhn',
-                        choices=['svhn'], 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')
-    
+        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"],
+        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",
+    )
+
+    # 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="cuda",
+        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()
-    
-
-    createfolders(args)
-    
-    device = 'cuda' if th.cuda.is_available() else 'cpu'
-    
-    #load model
-    model = load_model()
+
+    createfolders(args.datafolder, args.resultfolder, args.modelfolder)
+
+    device = args.device
+
+    # load model
+    model = load_model(args.modelname)
     model.to(device)
-    
+
     metrics = MetricWrapper(*args.metric)
-    
-    #Dataset
-    traindata = load_data(args.dataset)
-    validata = load_data(args.dataset)
-    
+
+    # Dataset
+    traindata = load_data(
+        args.dataset,
+        train=True,
+        data_path=args.datafolder,
+        download=args.download_data,
+    )
+    validata = load_data(
+        args.dataset,
+        train=False,
+        data_path=args.datafolder,
+    )
+
     trainloader = DataLoader(traindata,
                              batch_size=args.batchsize,
                              shuffle=True,
@@ -72,47 +152,50 @@ def main():
                             batch_size=args.batchsize,
                             shuffle=False,
                             pin_memory=True)
-    
+
     criterion = nn.CrossEntropyLoss()
-    optimizer = th.optim.Adam(model.parameters(), lr = args.learning_rate)
-    
-    
+    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")
+        exit(0)
+
     wandb.init(project='',
                tags=[])
     wandb.watch(model)
-    
+
     for epoch in range(args.epoch):
-        
-        #Training loop start
+
+        # Training loop start
         trainingloss = []
         model.train()
-        for x, y in traindata:
+        for x, y in trainloader:
             x, y = x.to(device), y.to(device)
             pred = model.forward(x)
-             
+
             loss = criterion(y, pred)
             loss.backward()
-            
+
             optimizer.step()
             optimizer.zero_grad(set_to_none=True)
             trainingloss.append(loss.item())
-        
+
         evalloss = []
-        #Eval loop start
+        # Eval loop start
         model.eval()
         with th.no_grad():
             for x, y in valiloader:
                 x = x.to(device)
                 pred = model.forward(x)
                 loss = criterion(y, pred)
                 evalloss.append(loss.item())
-                
+
         wandb.log({
             'Epoch': epoch,
             'Train loss': np.mean(trainingloss),
             'Evaluation Loss': np.mean(evalloss)
         })
-              
 
 
 if __name__ == '__main__':

utils/__init__.py

@@ -1,3 +1,5 @@
+__all__ = ['createfolders', 'load_data', 'load_model', 'MetricWrapper']
+
 from .createfolders import createfolders
 from .load_data import load_data
 from .load_metric import MetricWrapper

utils/createfolders.py

@@ -1,9 +1,9 @@
 import argparse
-import os
+from pathlib import Path
 from tempfile import TemporaryDirectory
 
 
-def createfolders(args) -> None:
+def createfolders(*dirs: Path) -> None:
     """
     Creates folders for storing data, results, model weights.
 
@@ -14,43 +14,44 @@ def createfolders(args) -> None:
 
     """
 
-    if not os.path.exists(args.datafolder):
-        os.makedirs(args.datafolder)
-        print(f"Created a folder at {args.datafolder}")
-
-    if not os.path.exists(args.resultfolder):
-        os.makedirs(args.resultfolder)
-        print(f"Created a folder at {args.resultfolder}")
-
-    if not os.path.exists(args.modelfolder):
-        os.makedirs(args.modelfolder)
-        print(f"Created a folder at {args.modelfolder}")
+    for dir in dirs:
+        dir.mkdir(parents=True, exist_ok=True)
 
 
 def test_createfolders():
-    with TemporaryDirectory(dir="tmp/") as temp_dir:
+    with TemporaryDirectory() as temp_dir:
+        temp_dir = Path(temp_dir)
+
         parser = argparse.ArgumentParser()
+
         # Structuture related values
         parser.add_argument(
             "--datafolder",
-            type=str,
-            default=os.path.join(temp_dir, "Data/"),
+            type=Path,
+            default=temp_dir / "Data",
             help="Path to where data will be saved during training.",
         )
         parser.add_argument(
             "--resultfolder",
-            type=str,
-            default=os.path.join(temp_dir, "Results/"),
+            type=Path,
+            default=temp_dir / "Results",
             help="Path to where results will be saved during evaluation.",
         )
         parser.add_argument(
             "--modelfolder",
-            type=str,
-            default=os.path.join(temp_dir, "Experiments/"),
+            type=Path,
+            default=temp_dir / "Experiments",
             help="Path to where model weights will be saved at the end of training.",
         )
 
-        args = parser.parse_args()
-        createfolders(args)
+        args = parser.parse_args([
+            "--datafolder", temp_dir / "Data",
+            "--resultfolder", temp_dir / "Results",
+            "--modelfolder", temp_dir / "Experiments"
+        ])
+
+        createfolders(args.datafolder, args.resultfolder, args.modelfolder)
 
-    return
+        assert (temp_dir / "Data").exists()
+        assert (temp_dir / "Results").exists()
+        assert (temp_dir / "Experiments").exists()

utils/dataloaders/__init__.py

@@ -0,0 +1,3 @@
+__all__ = ["USPSDataset0_6"]
+
+from .usps_0_6 import USPSDataset0_6