optimised

js_evaluator/JS_Evaluator.py

@@ -0,0 +1,111 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+import torchvision.models as models
+
+class ZeroRetrainForgettingEvaluator:
+    def __init__(self, unlearned_model: nn.Module, num_classes: int):
+        """
+        Initializes the ZRF Evaluator.
+        
+        Args:
+            unlearned_model (nn.Module): Your fine-tuned & unlearned ResNet-50.
+            num_classes (int): Number of classes used in your CelebA task.
+        """
+        # select device
+        if torch.cuda.is_available():
+            self.device = torch.device("cuda")
+        elif hasattr(torch, "xpu") and torch.xpu.is_available():
+            self.device = torch.device("xpu") # For Intel GPUs using IPEX
+        else:
+            self.device = torch.device("cpu")
+            
+        print(f"[INFO] Using device: {self.device}")
+
+        # prepare the unlearned model
+        self.unlearned_model = unlearned_model.to(self.device)
+        self.unlearned_model.eval()
+
+        # Instantiate a structurally matching, completely random model
+        print(f"[INFO] Initializing random baseline ResNet-50 with {num_classes} classes...")
+        self.random_model = self.get_random_model(num_classes)
+        self.random_model = self.random_model.to(self.device)
+        self.random_model.eval()
+
+   # gets randomly initialised model 
+   # for comparison with unlearned model
+    def get_random_model(num_classes):
+        print(f"[INFO] Initializing random baseline ResNet-50 with {num_classes} classes...")
+        model = models.resnet50(weights=None)
+        model.fc = nn.Linear(model.fc.in_features, num_classes)
+        return model
+
+
+    # compute divergence
+    def _compute_js_divergence(self, p: torch.Tensor, q: torch.Tensor) -> float:
+        """
+        Computes the Jensen-Shannon (JS) Divergence between two probability distributions.
+        
+        Args:
+            p, q (Tensor): Tensors of shape (batch_size, num_classes) containing probabilities.
+        """
+        # Avoid log(0) issues by adding a tiny epsilon
+        eps = 1e-12
+        p = torch.clamp(p, eps, 1.0)
+        q = torch.clamp(q, eps, 1.0)
+        
+        # Calculate the midpoint distribution
+        m = 0.5 * (p + q)
+        
+        # Compute KL Divergence natively: KL(P || M) and KL(Q || M)
+        kl_pm = torch.sum(p * (torch.log(p) - torch.log(m)), dim=1)
+        kl_qm = torch.sum(q * (torch.log(q) - torch.log(m)), dim=1)
+        
+        # JS Divergence is the average of both KL divergences
+        js_div = 0.5 * (kl_pm + kl_qm)
+        
+        # Return the mean divergence across the entire batch
+        return js_div.mean().item()
+
+    def evaluate_forget_class(self, dataset, batch_size: int = 32) -> float:
+        """
+        Evaluates the unlearned model against the random model using images 
+        from the forgotten class/identity.
+        
+        Args:
+            dataset (Dataset): A PyTorch Dataset containing images of the forget set.
+            batch_size (int): Batch size for evaluation.
+            
+        Returns:
+            float: The ZRF score (JS Divergence). A lower divergence means 
+                   the unlearned model is behaving exactly like a random model.
+        """
+        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
+        
+        total_js_div = 0.0
+        total_samples = 0
+
+        # No gradients needed for evaluation
+        with torch.no_grad():
+            for images, _ in dataloader:
+                images = images.to(self.device)
+                batch_len = images.size(0)
+
+                # Get raw outputs (logits)
+                unlearned_logits = self.unlearned_model(images)
+                random_logits = self.random_model(images)
+
+                # Convert logits to probability distributions via Softmax
+                unlearned_probs = F.softmax(unlearned_logits, dim=1)
+                random_probs = F.softmax(random_logits, dim=1)
+
+                # Calculate JS divergence for this batch
+                batch_js = self._compute_js_divergence(unlearned_probs, random_probs)
+                
+                # Weighted average based on batch size (handles final smaller batches perfectly)
+                total_js_div += batch_js * batch_len
+                total_samples += batch_len
+
+        final_zrf_score = total_js_div / total_samples
+        return final_zrf_score