added evaulations

models/training_environment.py

@@ -1,10 +1,16 @@
 import importlib
-from models.utils import calculate_metrics
-
+from models.utils import calculate_metrics, plot_to_image, get_dance_mapping
+import numpy as np
 from abc import ABC, abstractmethod
 import pytorch_lightning as pl
+import matplotlib.pyplot as plt
 import torch
 import torch.nn as nn
+from sklearn.metrics import (
+    roc_auc_score,
+    confusion_matrix,
+    ConfusionMatrixDisplay,
+)
 
 
 class TrainingEnvironment(pl.LightningModule):
@@ -27,8 +33,8 @@ class TrainingEnvironment(pl.LightningModule):
             config["training_environment"].get("loggers", {})
         )
         self.config = config
-        self.has_multi_label_predictions = (
-            not type(criterion).__name__ == "CrossEntropyLoss"
+        self.has_multi_label_predictions = not (
+            type(criterion).__name__ == "CrossEntropyLoss"
         )
         self.save_hyperparameters(
             {
@@ -44,6 +50,8 @@ class TrainingEnvironment(pl.LightningModule):
     ) -> torch.Tensor:
         features, labels = batch
         outputs = self.model(features)
+        if self.has_multi_label_predictions:
+            outputs = nn.functional.sigmoid(outputs)
         loss = self.criterion(outputs, labels)
         metrics = calculate_metrics(
             outputs,
@@ -62,6 +70,8 @@ class TrainingEnvironment(pl.LightningModule):
     ):
         x, y = batch
         preds = self.model(x)
+        if self.has_multi_label_predictions:
+            preds = nn.functional.sigmoid(preds)
         metrics = calculate_metrics(
             preds, y, prefix="val/", multi_label=self.has_multi_label_predictions
         )
@@ -71,12 +81,48 @@ class TrainingEnvironment(pl.LightningModule):
     def test_step(self, batch: tuple[torch.Tensor, torch.TensorType], batch_index: int):
         x, y = batch
         preds = self.model(x)
-        self.log_dict(
-            calculate_metrics(
-                preds, y, prefix="test/", multi_label=self.has_multi_label_predictions
-            ),
-            prog_bar=True,
+        if self.has_multi_label_predictions:
+            preds = nn.functional.sigmoid(preds)
+        metrics = calculate_metrics(
+            preds, y, prefix="test/", multi_label=self.has_multi_label_predictions
+        )
+        if not self.has_multi_label_predictions:
+            preds = nn.functional.softmax(preds, dim=1)
+        y = y.detach().cpu().numpy()
+        preds = preds.detach().cpu().numpy()
+        # ROC-auc score
+        try:
+            metrics["test/roc_auc_score"] = torch.tensor(
+                roc_auc_score(y, preds), dtype=torch.float32
+            )
+        except ValueError:
+            # If there is only one class, roc_auc_score will throw an error
+            pass
+
+            pass
+        self.log_dict(metrics, prog_bar=True)
+        # Create confusion matrix
+
+        preds = preds.argmax(axis=1)
+        y = y.argmax(axis=1)
+        cm = confusion_matrix(
+            preds, y, normalize="all", labels=np.arange(len(self.config["dance_ids"]))
         )
+        if hasattr(self, "test_cm"):
+            self.test_cm += cm
+        else:
+            self.test_cm = cm
+
+    def on_test_end(self):
+        dance_ids = sorted(self.config["dance_ids"])
+        np.fill_diagonal(self.test_cm, 0)
+        cm = self.test_cm / self.test_cm.max()
+        ConfusionMatrixDisplay(cm, display_labels=dance_ids).plot()
+        image = plot_to_image(plt.gcf())
+        image = torch.tensor(image, dtype=torch.uint8)
+        image = image.permute(2, 0, 1)
+        self.logger.experiment.add_image("test/confusion_matrix", image, 0)
+        delattr(self, "test_cm")
 
     def configure_optimizers(self):
         optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)

models/utils.py

@@ -2,6 +2,11 @@ import torch.nn as nn
 import torch
 import numpy as np
 from sklearn.metrics import precision_score, recall_score, f1_score, accuracy_score
+from functools import cache
+import pandas as pd
+import matplotlib.pyplot as plt
+import io
+from PIL import Image
 
 
 class LabelWeightedBCELoss(nn.Module):
@@ -38,10 +43,13 @@ def calculate_metrics(
 ) -> dict[str, torch.Tensor]:
     target = target.detach().cpu().numpy()
     pred = pred.detach().cpu()
-    pred = nn.functional.softmax(pred, dim=1)
+    if not multi_label:
+        pred = nn.functional.softmax(pred, dim=1)
     pred = pred.numpy()
     params = {
-        "y_true": target if multi_label else target.argmax(1),
+        "y_true": np.array(target > 0.0, dtype=float)
+        if multi_label
+        else target.argmax(1),
         "y_pred": np.array(pred > threshold, dtype=float)
         if multi_label
         else pred.argmax(1),
@@ -85,3 +93,22 @@ def get_id_label_mapping(labels: list[str]) -> tuple[dict, dict]:
 def compute_hf_metrics(eval_pred):
     predictions = np.argmax(eval_pred.predictions, axis=1)
     return accuracy_score(y_true=eval_pred.label_ids, y_pred=predictions)
+
+
+@cache
+def get_dance_mapping(mapping_file: str) -> dict[str, str]:
+    mapping_df = pd.read_csv(mapping_file)
+    return {row["id"]: row["name"] for _, row in mapping_df.iterrows()}
+
+
+def plot_to_image(figure) -> np.ndarray:
+    """Converts the matplotlib plot specified by 'figure' to a PNG image and
+    returns it. The supplied figure is closed and inaccessible after this call."""
+    # Save the plot to a PNG in memory.
+    buf = io.BytesIO()
+    plt.savefig(buf, format="png")
+    # Closing the figure prevents it from being displayed directly inside
+    # the notebook.
+    plt.close(figure)
+    buf.seek(0)
+    return np.array(Image.open(buf))