added num_classes parameter to Metric wrapper and metric calculation step to train, val loops

Author: hzavadil98

main.py

@@ -31,7 +31,6 @@ def main():
 
     device = args.device
 
-    metrics = MetricWrapper(*args.metric)
 
     if args.dataset.lower() == "usps_0-6" or args.dataset.lower() == "uspsh5_7_9":
         augmentations = transforms.Compose(
@@ -59,6 +58,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:
@@ -90,24 +91,28 @@ 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)
+            
+            preds = th.argmax(logits, dim=1)
+            metrics(y, preds)
 
-            break
 
+            break
+        print(metrics.__getmetrics__())
         print("Dry run completed successfully.")
         exit(0)
 
@@ -120,24 +125,36 @@ 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(y, logits)
                 evalloss.append(loss.item())
+                
+                preds = th.argmax(logits, dim=1)
+                metrics(y, preds)
+        
+        wandb.log(metrics.__getmetrics__(str_prefix="Evaluation "))
+        metrics.__resetvalues__()
 
         wandb.log(
             {

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
@@ -149,3 +149,5 @@ def __getitem__(self, index):
             image = self.transform(image)
 
         return image, label
+    
+

utils/load_metric.py

@@ -3,13 +3,14 @@
 import numpy as np
 import torch.nn as nn
 
-from .metrics import Accuracy, EntropyPrediction, F1Score, Precision
+from .metrics import Accuracy, EntropyPrediction, F1Score, Precision, Recall
 
 
 class MetricWrapper(nn.Module):
-    def __init__(self, *metrics):
+    def __init__(self, *metrics, num_classes):
         super().__init__()
         self.metrics = {}
+        self.num_classes = num_classes
 
         for metric in metrics:
             self.metrics[metric] = self._get_metric(metric)
@@ -33,26 +34,29 @@ def _get_metric(self, key):
 
         match key.lower():
             case "entropy":
-                return EntropyPrediction()
+                return EntropyPrediction(num_classes=self.num_classes)
             case "f1":
-                raise F1Score()
+                return F1Score(num_classes=self.num_classes)
             case "recall":
-                raise NotImplementedError("Recall score not implemented yet")
+                return Recall(num_classes=self.num_classes)
             case "precision":
-                return Precision()
+                return Precision(num_classes=self.num_classes)
             case "accuracy":
-                return Accuracy()
+                return Accuracy(num_classes=self.num_classes)
             case _:
                 raise ValueError(f"Metric {key} not supported")
 
     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):
+    def __getmetrics__(self, str_prefix: str = None):
         return_metrics = {}
         for key in self.metrics:
-            return_metrics[key] = np.mean(self.tmp_scores[key])
+            if str_prefix is not None:
+                return_metrics[str_prefix + key] = np.mean(self.tmp_scores[key])
+            else:
+                return_metrics[key] = np.mean(self.tmp_scores[key])
 
         return return_metrics
 

utils/metrics/accuracy.py

@@ -3,9 +3,10 @@
 
 
 class Accuracy(nn.Module):
-    def __init__(self):
+    def __init__(self, num_classes):
         super().__init__()
-
+        self.num_classes = num_classes
+        
     def forward(self, y_true, y_pred):
         """
         Compute the accuracy of the model.