optimised

Predict.py

@@ -0,0 +1,67 @@
+
+
+
+import torch
+import numpy as np
+
+@torch.inference_mode() # More memory-efficient than no_grad()
+def get_loss_per_sample(model, data_loader, device):
+    """
+    Returns a list of individual losses for every sample in the loader.
+    Useful for MIA to see how 'certain' the model is about specific images.
+    """
+    model.eval()
+    criterion = torch.nn.CrossEntropyLoss(reduction='none') # Crucial: returns loss per image
+    all_losses = []
+
+    for inputs, labels in data_loader:
+        inputs, labels = inputs.to(device), labels.to(device)
+        
+        outputs = model(inputs)
+        
+        # Calculate loss for each image in the batch individually
+        loss = criterion(outputs, labels)
+        
+        all_losses.extend(loss.cpu().numpy())
+
+    return all_losses
+
+
+@torch.inference_mode()
+def get_losses_by_class(model, data_loader, device):
+    """
+    Returns a dictionary: { class_id: [list_of_losses_for_this_class] }
+    """
+    model.eval()
+    criterion = torch.nn.CrossEntropyLoss(reduction='none')
+    
+    class_losses = {}
+
+    for inputs, labels in data_loader:
+        inputs, labels = inputs.to(device), labels.to(device)
+        outputs = model(inputs)
+        
+        # Get individual losses
+        losses = criterion(outputs, labels).cpu().numpy()
+        labels_np = labels.cpu().numpy()
+
+        for i, class_id in enumerate(labels_np):
+            if class_id not in class_losses:
+                class_losses[class_id] = []
+            class_losses[class_id].append(losses[i])
+
+    return class_losses
+
+
+# evaluate MIA
+def eval_MIA(forgotten_losses, never_seen_losses):
+    avg_f_loss = np.mean(forgotten_losses)
+    avg_ns_loss = np.mean(never_seen_losses)
+
+    print(f"Average Loss on Forgotten Identity: {avg_f_loss:.4f}")
+    print(f"Average Loss on Unknown Identities: {avg_ns_loss:.4f}")
+
+    if avg_f_loss < avg_ns_loss * 0.8:
+        print("MIA Warning: Model still shows high certainty on forgotten data.")
+    else:
+        print("MIA Success: Model treats forgotten data like unknown data.")