Merge pull request #16 from SFI-Visual-Intelligence/magnus-branch

Changed metric loader to a wrapper

Author: Seilmast

main.py

@@ -1,18 +1,10 @@
 import torch as th
 import torch.nn as nn 
-
+from torch.utils.data import DataLoader
 import argparse
-from utils import load_metric, load_model, createfolders
-
-
-
-
-
-
-
-
-
-
+import wandb
+import numpy as np
+from utils import MetricWrapper, load_model, load_data, createfolders
 
 
 def main():
@@ -38,16 +30,28 @@ def main():
     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('--EntropyPrediction', type=bool, default=True, help='Include the Entropy Prediction metric in evaluation')
+    parser.add_argument('--F1Score', type=bool, default=True, help='Include the F1Score metric in evaluation')
+    parser.add_argument('--Recall', type=bool, default=True, help='Include the Recall metric in evaluation')
+    parser.add_argument('--Precision', type=bool, default=True, help='Include the Precision metric in evaluation')
+    parser.add_argument('--Accuracy', type=bool, default=True, help='Include the Accuracy metric in 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')
+    
     args = parser.parse_args()
+    
 
     createfolders(args)
 
@@ -57,19 +61,68 @@ def main():
     model = load_model()
     model.to(device)
 
+    metrics = MetricWrapper(
+        EntropyPred = args.EntropyPrediction,
+        F1Score = args.F1Score,
+        Recall = args.Recall,
+        Precision = args.Precision,
+        Accuracy = args.Accuracy
+    )
+    
+    #Dataset
+    traindata = load_data(args.dataset)
+    validata = load_data(args.dataset)
+    
+    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)
-
-
-
+    
+    
+    wandb.init(project='',
+               tags=[])
+    wandb.watch(model)
+    
     for epoch in range(args.epoch):
 
         #Training loop start
+        trainingloss = []
+        model.train()
+        for x, y in traindata:
+            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
-    
-        pass 
+        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__':

test.ipynb

@@ -0,0 +1,76 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import argparse"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "usage: [-h] [--datafolder DATAFOLDER]\n",
+      ": error: unrecognized arguments: --f=/home/magnus/.local/share/jupyter/runtime/kernel-v3fc3d3b04bd8d83becf1be5eacf19e7bf46887012.json\n"
+     ]
+    },
+    {
+     "ename": "SystemExit",
+     "evalue": "2",
+     "output_type": "error",
+     "traceback": [
+      "An exception has occurred, use %tb to see the full traceback.\n",
+      "\u001b[0;31mSystemExit\u001b[0m\u001b[0;31m:\u001b[0m 2\n"
+     ]
+    }
+   ],
+   "source": [
+    "parser = argparse.ArgumentParser(\n",
+    "        prog='',\n",
+    "        description='',\n",
+    "        epilog='',\n",
+    "    )\n",
+    "parser.add_argument('--datafolder', type=str, default='Data/', help='Path to where data will be saved during training.')\n",
+    "args = parser.parse_args()\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(args)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "env",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

utils/__init__.py

@@ -1,4 +1,4 @@
-from .load_metric import load_metric
+from .load_metric import MetricWrapper
 from .load_model import load_model
 from .load_data import load_data
 from .createfolders import createfolders

utils/load_metric.py

@@ -1,27 +1,49 @@
+import copy 
+import numpy as np 
 import torch.nn as nn 
 from metrics import EntropyPrediction
 
-def load_metric(metricname:str) -> nn.Module:
-    '''
-    This function returns an instance of a class inhereting from nn.Module. 
-    This class returns the given metric given a set of label - prediction pairs. 
-    
-    Parameters
-    ----------
-    metricname: string
-        string naming the metric to return. 
-    
-    Returns
-    -------
-    Class
-        Returns an instance of a class inhereting from nn.Module.
-    
-    Raises
-    ------
-    ValueError
-        When the metricname parameter does not match any implemented metric, raise this error along with a descriptive message. 
-    '''
-    if metricname == 'EntropyPrediction':
-        return EntropyPrediction()
-    else:
-        raise ValueError(f'Metric: {metricname} has not been implemented. \nCheck the documentation for implemented metrics, or check your spelling')
+
+class MetricWrapper(nn.Module):
+    def __init__(self,
+                 EntropyPred:bool = True, 
+                 F1Score:bool = True,
+                 Recall:bool = True,
+                 Precision:bool = True,
+                 Accuracy:bool = True):
+        super().__init__()
+        self.metrics = {}
+        
+        if EntropyPred:
+            self.metrics['Entropy of Predictions'] = EntropyPrediction()
+
+        if F1Score:
+            self.metrics['F1 Score'] = None 
+        
+        if Recall:
+            self.metrics['Recall'] = None 
+        
+        if Precision:
+            self.metrics['Precision'] = None
+        
+        if Accuracy:
+            self.metrics['Accuracy'] = None
+            
+        self.tmp_scores = copy.deepcopy(self.metrics)
+        for key in self.tmp_scores:
+            self.tmp_scores[key] = []
+
+    def __call__(self, y_true, y_pred):
+        for key in self.metrics:
+            self.tmp_scores[key].append(self.metrics[key](y_true, y_pred))
+        
+    def __getmetrics__(self):
+        return_metrics = {}
+        for key in self.metrics:
+            return_metrics[key] = np.mean(self.tmp_scores[key])
+        
+        return return_metrics
+
+    def __resetvalues__(self):
+        for key in self.tmp_scores:
+            self.tmp_scores[key] = []