Fixed thing

Author: Seilmast

main.py

@@ -107,8 +107,7 @@ def main():
             optimizer.step()
             optimizer.zero_grad(set_to_none=True)
 
-            preds = th.argmax(logits, dim=1)
-            metrics(y, preds)
+            metrics(y, logits)
 
             break
         print(metrics.accumulate())
@@ -134,8 +133,7 @@ def main():
             optimizer.zero_grad(set_to_none=True)
             trainingloss.append(loss.item())
 
-            preds = th.argmax(logits, dim=1)
-            metrics(y, preds)
+            metrics(y, logits)
 
         wandb.log(metrics.accumulate(str_prefix="Train "))
         metrics.reset()
@@ -150,8 +148,7 @@ def main():
                 loss = criterion(logits, y)
                 evalloss.append(loss.item())
 
-                preds = th.argmax(logits, dim=1)
-                metrics(y, preds)
+                metrics(y, logits)
 
         wandb.log(metrics.accumulate(str_prefix="Evaluation "))
         metrics.reset()

utils/metrics/EntropyPred.py

@@ -1,12 +1,32 @@
+import numpy as np
 import torch.nn as nn
+from scipy.stats import entropy
 
 
 class EntropyPrediction(nn.Module):
-    def __init__(self):
+    def __init__(self, averages: str = 'average'):
+        """
+        Initializes the EntropyPrediction module.
+        Args:
+            averages (str): Specifies the method of aggregation for entropy values. 
+                            Must be either 'average' or 'sum'.
+        Raises:
+            AssertionError: If the averages parameter is not 'average' or 'sum'.
+        """
         super().__init__()
-
+        
+        assert averages == 'average' or averages == 'sum'
+        self.averages = averages
+        self.stored_entropy_values = []
+        
     def __call__(self, y_true, y_false_logits):
-        return
-
-    def __reset__(self):
-        pass
+        """
+        Computes the entropy between true labels and predicted logits, storing the results.
+        Args:
+            y_true: The true labels.
+            y_false_logits: The predicted logits.
+        Side Effects:
+            Appends the computed entropy values to the stored_entropy_values list.
+        """
+        entropy_values = entropy(y_true, qk=y_false_logits)
+        return entropy_values