Added computation of 95% bootstrap confidence intervals

Author: PierreMarza

src/thunder/tasks/adversarial_attack.py

@@ -15,13 +15,19 @@
 import torchvision.transforms as T
 import wandb
 from omegaconf import DictConfig
+from sklearn.metrics import (
+    accuracy_score,
+    balanced_accuracy_score,
+    f1_score,
+    jaccard_score,
+)
 from torch.utils.data import DataLoader, Subset
 from tqdm import tqdm
 
 from ..models.pretrained_models import load_pretrained_model
 from ..utils.constants import UtilsConstants
 from ..utils.data import PatchDataset, get_data
-from ..utils.downstream_metrics import compute_metrics
+from ..utils.downstream_metrics import compute_metric, compute_metrics
 from ..utils.pgd_attack_linear import PGDImageAttack
 from ..utils.utils import log_metrics, save_outputs
 
@@ -177,14 +183,51 @@ def adversarial_attack(
     all_adv_preds = torch.cat(all_adv_preds).cpu().numpy()
     all_labels = torch.cat(all_labels).cpu().numpy()
 
-    # Compute accuracy before and after attack
+    # Compute metrics before and after attack
     clean_metrics = compute_metrics(None, all_clean_preds, all_labels)
     adv_metrics = compute_metrics(None, all_adv_preds, all_labels)
     metrics = {
         "clean": clean_metrics,
         "adversarial": adv_metrics,
     }
 
+    # Drop in metrics (mean and confidence intervals)
+    f1_drop = compute_metric(
+        all_labels,
+        np.concatenate([all_clean_preds[:, None], all_adv_preds[:, None]], axis=1),
+        lambda y, y_pred: f1_score(y_true=y, y_pred=y_pred[:, 0], average="macro")
+        - f1_score(y_true=y, y_pred=y_pred[:, 1], average="macro"),
+        label_indices=np.arange(len(all_labels)),
+    )
+    accuracy_drop = compute_metric(
+        all_labels,
+        np.concatenate([all_clean_preds[:, None], all_adv_preds[:, None]], axis=1),
+        lambda y, y_pred: accuracy_score(y_true=y, y_pred=y_pred[:, 0])
+        - accuracy_score(y_true=y, y_pred=y_pred[:, 1]),
+        label_indices=np.arange(len(all_labels)),
+    )
+    jaccard_drop = compute_metric(
+        all_labels,
+        np.concatenate([all_clean_preds[:, None], all_adv_preds[:, None]], axis=1),
+        lambda y, y_pred: jaccard_score(y_true=y, y_pred=y_pred[:, 0], average="macro")
+        - jaccard_score(y_true=y, y_pred=y_pred[:, 1], average="macro"),
+        label_indices=np.arange(len(all_labels)),
+    )
+    balanced_accuracy_drop = compute_metric(
+        all_labels,
+        np.concatenate([all_clean_preds[:, None], all_adv_preds[:, None]], axis=1),
+        lambda y, y_pred: balanced_accuracy_score(y_true=y, y_pred=y_pred[:, 0])
+        - balanced_accuracy_score(y_true=y, y_pred=y_pred[:, 1]),
+        label_indices=np.arange(len(all_labels)),
+    )
+
+    metrics["drop"] = {
+        "f1": f1_drop,
+        "accuracy": accuracy_drop,
+        "jaccard": jaccard_drop,
+        "balanced_accuracy": balanced_accuracy_drop,
+    }
+
     # save ---------------------------------------------------------------
     save_outputs(res_folder, metrics)
 

src/thunder/tasks/image_retrieval.py

@@ -35,6 +35,7 @@ def topk_retrieval(
     chunk_size: int = 10000,
     return_viz_data: bool = False,
     disable_progress_bar: bool = False,
+    compute_ci: bool = True,
 ) -> tuple[dict, dict, dict]:
     """
     Computing similarities between queries and keys.
@@ -46,6 +47,7 @@ def topk_retrieval(
     :param chunk_size: maximum number of query embeddings for which we compute dot product similarity with all key embeddings.
     :param return_viz_data: whether to return data to visualize topk samples.
     :param disable_progress_bar: whether to hide the progress bar.
+    :param compute_ci: whether to compute confidence intervals.
     :return dict of metrics, sorted images ids and viz data if required.
     """
     # Normalizing embeddings
@@ -93,7 +95,9 @@ def topk_retrieval(
     # Metrics
     metrics_per_k = {}
     for k in k_vals:
-        metrics_per_k[k] = compute_metrics(None, np.array(preds_per_k[k]), query_labels)
+        metrics_per_k[k] = compute_metrics(
+            None, np.array(preds_per_k[k]), query_labels, compute_ci=compute_ci
+        )
 
     return metrics_per_k, sorted_ids_per_k, viz_data
 
@@ -180,7 +184,7 @@ def image_retrieval(
 
     # Logging
     for k in k_vals:
-        log_metrics(wandb_base_folder, metrics, "test", step=k)
+        log_metrics(wandb_base_folder, metrics[k], "test", step=k)
     save_outputs(res_folder, metrics)
 
     return metrics

src/thunder/tasks/knn_classification.py

@@ -30,6 +30,7 @@ def knn(
         embs["val"],
         labels["val"],
         k_vals,
+        compute_ci=False,
     )
 
     # Logging
@@ -40,7 +41,7 @@ def knn(
     best_k = None
     best_val_f1 = -float("inf")
     for k in k_vals:
-        val_f1 = np.array(val_metrics[k]["f1"]).mean().item()
+        val_f1 = np.array(val_metrics[k]["f1"]["metric_score"]).mean().item()
         if val_f1 > best_val_f1:
             best_val_f1 = val_f1
             best_k = k

src/thunder/tasks/simple_shot.py

@@ -77,6 +77,7 @@ def simple_shot(
                 test_labels,
                 [1],
                 disable_progress_bar=True,
+                compute_ci=False,
             )
 
             # Logging

src/thunder/tasks/train_eval_probe.py

@@ -26,7 +26,7 @@
 from ..utils.constants import UtilsConstants
 from ..utils.calibration_metrics import compute_calibration_metrics
 from ..utils.data import PatchDataset
-from ..utils.downstream_metrics import compute_metrics
+from ..utils.downstream_metrics import compute_metric, compute_metrics
 from ..utils.utils import (
     get_hyperaparams_dict,
     local_seed,
@@ -272,13 +272,13 @@ def train_probe(
                     # Updating best ckpt
                     if (
                         task_type == "linear_probing"
-                        and metrics[i]["f1"] > best_val_perf
+                        and metrics[i]["f1"]["metric_score"] > best_val_perf
                     ) or (
                         task_type == "segmentation"
                         and np.array(losses[i]).mean().item() < best_val_perf
                     ):
                         if task_type == "linear_probing":
-                            best_val_perf = metrics[i]["f1"]
+                            best_val_perf = metrics[i]["f1"]["metric_score"]
                         elif task_type == "segmentation":
                             best_val_perf = np.array(losses[i]).mean().item()
                         best_ckpt_hyperparam_id = i
@@ -538,7 +538,9 @@ def train_eval(
     tot_loss = []
     all_out = []
     all_label = []
-    for batch_id, batch in tqdm(enumerate(dataloader), total=len(dataloader)):
+    for batch_id, batch in tqdm(
+        enumerate(dataloader), total=len(dataloader), disable=hyperparam_search
+    ):
         # Batch data
         if "emb" in batch.keys():
             emb = batch["emb"].to(device)
@@ -591,8 +593,14 @@ def train_eval(
             out = []
             if task_type == "segmentation":
                 for o, m in zip(output, unmasked_label):
-                    out.append(torch.cat([o[c][m].unsqueeze(-1) for c in range(o.shape[0])], dim=-1))
-                curr_loss = sum([criterion(o, l) for o, l in zip(out, label)]) / len(out)
+                    out.append(
+                        torch.cat(
+                            [o[c][m].unsqueeze(-1) for c in range(o.shape[0])], dim=-1
+                        )
+                    )
+                curr_loss = sum([criterion(o, l) for o, l in zip(out, label)]) / len(
+                    out
+                )
             else:
                 for c in range(output.shape[1]):
                     out.append(output[:, c].unsqueeze(-1))
@@ -604,13 +612,19 @@ def train_eval(
             if batch_id == 0:
                 tot_loss.append([curr_loss.item()])
                 if comp_metrics:
-                    all_out.append([[o.detach().cpu() for o in out]] if task_type == 'segmentation'
-                                   else [out.detach().cpu()])
+                    all_out.append(
+                        [[o.detach().cpu() for o in out]]
+                        if task_type == "segmentation"
+                        else [out.detach().cpu()]
+                    )
             else:
                 tot_loss[i].append(curr_loss.item())
                 if comp_metrics:
-                    all_out[i].append([o.detach().cpu() for o in out] if task_type == 'segmentation'
-                                      else out.detach().cpu())
+                    all_out[i].append(
+                        [o.detach().cpu() for o in out]
+                        if task_type == "segmentation"
+                        else out.detach().cpu()
+                    )
         # Logging
         if comp_metrics:
             if task_type == "segmentation":
@@ -641,20 +655,42 @@ def train_eval(
         if task_type == "segmentation":
             metrics = []
             for i in range(len(all_out)):
-                all_out[i] = [F.softmax(item, dim=1) for batch in all_out[i] for item in batch]
-                all_metrics = [compute_metrics(o, None, l, True) for o, l in zip(all_out[i], all_label) if len(l) > 0]
+                all_out[i] = [
+                    F.softmax(item, dim=1) for batch in all_out[i] for item in batch
+                ]
+                all_metrics = [
+                    compute_metrics(o, None, l, True, compute_ci=False)
+                    for o, l in zip(all_out[i], all_label)
+                    if len(l) > 0
+                ]
                 weights = [len(l) for l in all_label if len(l) > 0]
-                metrics.append({key: np.average([d[key] for d in all_metrics], weights=weights) for key in all_metrics[0]} |
-                               {f'{key}_per_sample': [d[key] for d in all_metrics] for key in all_metrics[0]})
+
+                # Averagin per-image performance and computing confidence intervals
+                all_metrics_out = {}
+                for key in all_metrics[0]:
+                    metric_vals = [d[key]["metric_score"] for d in all_metrics]
+                    all_metrics_out[key] = compute_metric(
+                        weights,
+                        metric_vals,
+                        lambda weights, metric_vals: np.average(
+                            metric_vals, weights=weights
+                        ),
+                    )
+                    all_metrics_out[f"per_sample_{key}"] = metric_vals
+                metrics.append(all_metrics_out)
         else:
             # Computing metrics
             all_label = torch.cat(all_label)
             metrics = []
             for i in range(len(all_out)):
                 all_out[i] = torch.cat(all_out[i])
                 all_out[i] = F.softmax(all_out[i], dim=1)
-                classification_metrics = compute_metrics(all_out[i], None, all_label)
-                conformal_metrics = compute_calibration_metrics(all_out[i], all_label)
+                classification_metrics = compute_metrics(
+                    all_out[i], None, all_label, compute_ci=(not hyperparam_search)
+                )
+                conformal_metrics = compute_calibration_metrics(
+                    all_out[i], all_label, compute_ci=(not hyperparam_search)
+                )
                 curr_metrics = (
                     classification_metrics | conformal_metrics
                 )  # merging dictionaries

src/thunder/utils/calibration_metrics.py

@@ -1,12 +1,17 @@
 import numpy as np
 import torch
 
+from ..utils.downstream_metrics import compute_metric
 
-def compute_calibration_metrics(out_proba: torch.Tensor, label: torch.Tensor) -> dict:
+
+def compute_calibration_metrics(
+    out_proba: torch.Tensor, label: torch.Tensor, compute_ci: bool = True
+) -> dict:
     """
     Computing performance metrics.
     :param out: tensor of proba predictions.
     :param label: tensor of ground-truth labels.
+    :param compute_ci: whether to compute confidence intervals.
     :return: dict of conformal prediction metrics.
     """
 
@@ -15,11 +20,41 @@ def compute_calibration_metrics(out_proba: torch.Tensor, label: torch.Tensor) ->
     label = label.numpy()
 
     # Computing metrics
-    ece = expected_calibration_error(out_proba, label)
-    mce = maximum_calibration_error(out_proba, label)
-    sce = static_calibration_error(out_proba, label)
-    ace = adaptive_calibration_error(out_proba, label)
-    tace = thresholded_adaptive_calibration_error(out_proba, label)
+    ece = compute_metric(
+        label,
+        out_proba,
+        lambda y, y_proba: expected_calibration_error(y_proba, y),
+        label_indices=np.arange(len(label)),
+        compute_ci=compute_ci,
+    )
+    mce = compute_metric(
+        label,
+        out_proba,
+        lambda y, y_proba: maximum_calibration_error(y_proba, y),
+        label_indices=np.arange(len(label)),
+        compute_ci=compute_ci,
+    )
+    sce = compute_metric(
+        label,
+        out_proba,
+        lambda y, y_proba: static_calibration_error(y_proba, y),
+        label_indices=np.arange(len(label)),
+        compute_ci=compute_ci,
+    )
+    ace = compute_metric(
+        label,
+        out_proba,
+        lambda y, y_proba: adaptive_calibration_error(y_proba, y),
+        label_indices=np.arange(len(label)),
+        compute_ci=compute_ci,
+    )
+    tace = compute_metric(
+        label,
+        out_proba,
+        lambda y, y_proba: thresholded_adaptive_calibration_error(y_proba, y),
+        label_indices=np.arange(len(label)),
+        compute_ci=compute_ci,
+    )
 
     return {
         "ECE": ece,