attck metrics

eval/UnlearningAttack.py

@@ -0,0 +1,238 @@
+import torch
+import torch.nn as nn
+import numpy as np
+import os
+from torch.utils.data import DataLoader
+from sklearn.linear_model import LogisticRegression
+from sklearn.metrics import accuracy_score
+from architectures.Model import Model
+
+class UnlearningAttack:
+    def __init__(self, arch, class_size):
+        """
+        Initializes the robust verification suite with universal architecture metadata.
+        Matches Section 5.5 of the thesis text by implementing distinct 
+        Parameter-Space and Logit-Space (Look-alike) attack pipelines uniformly.
+        """
+        self.arch = arch
+        self.class_size = class_size
+        
+        self.hook = None
+        self.model = None  
+        self._hook_features = []
+        self.criterion = nn.CrossEntropyLoss(reduction='none')
+        self.collecting = False
+
+    def _hook_fn(self, module, input, output):
+        if not self.collecting:
+            return
+        flattened_embeddings = torch.flatten(output, 1)
+        self._hook_features.append(flattened_embeddings.detach())
+    
+    def register_model_hook(self, inner_model):
+        if hasattr(inner_model, "original_model"):
+            core_model = inner_model.original_model
+        else:
+            core_model = inner_model
+
+        if hasattr(core_model, 'avgpool'):
+            pool_layer = core_model.avgpool
+        else:
+            pool_layer = None
+            for name, module in core_model.named_modules():
+                if 'pool' in name:
+                    pool_layer = module
+                    break
+
+        if pool_layer is None:
+            raise AttributeError("The target model architecture lacks an 'avgpool' layer block.")
+            
+        self.hook = pool_layer.register_forward_hook(self._hook_fn)
+
+    def shutdown_hook(self):
+        if hasattr(self, 'hook') and self.hook:
+            self.hook.remove()
+            self.hook = None
+
+    def _extract_attack_features(self, target_model, loader, device, target_class):
+        target_model.eval()
+        all_probs = []
+        all_entropies = []
+        all_losses = []
+        self._hook_features = [] 
+        self.collecting = True
+        with torch.no_grad():
+            for data, targets in loader:
+                data, targets = data.to(device), targets.to(device)
+                if target_model.__class__.__name__ == "WF_Module":
+                    gate_signals = torch.full(
+                        (data.size(0),), 
+                        target_class,  
+                        dtype=torch.long, 
+                        device=data.device
+                    )
+                    outputs = target_model(data, target_class_indices=gate_signals)
+                else:
+                    outputs = target_model(data)
+                
+                probs = torch.softmax(outputs, dim=1)
+                all_probs.extend(probs.cpu().numpy())
+                
+                entropy = -torch.sum(probs * torch.log(probs + 1e-10), dim=1)
+                all_entropies.extend(entropy.cpu().numpy())
+                
+                loss = self.criterion(outputs, targets)
+                all_losses.extend(loss.cpu().numpy())
+
+        self.collecting = False
+
+        X_probs = np.array(all_probs)
+        X_entropy = np.array(all_entropies).reshape(-1, 1)
+        X_loss = np.array(all_losses).reshape(-1, 1)
+        X_features = np.hstack([X_probs, X_entropy, X_loss])
+        
+        if self._hook_features:
+            compiled_latent = torch.cat(self._hook_features, dim=0).cpu().numpy()
+        else:
+            compiled_latent = np.zeros((len(X_features), 512))
+
+        return X_features, compiled_latent
+
+    def run_parameter_space_mia(self, unlearned_model, shadow_model, forget_loader, retain_test_loader, device, index):
+        X_shadow_mem, _ = self._extract_attack_features(shadow_model, forget_loader, device, index)
+        X_shadow_non, _ = self._extract_attack_features(shadow_model, retain_test_loader, device, index)
+        
+        min_train = min(len(X_shadow_mem), len(X_shadow_non))
+        X_train = np.vstack([X_shadow_mem[:min_train], X_shadow_non[:min_train]])
+        y_train = np.concatenate([np.ones(min_train), np.zeros(min_train)])
+        
+        attack_classifier = LogisticRegression(max_iter=1000)
+        attack_classifier.fit(X_train, y_train)
+        
+        X_eval_mem, latent_features = self._extract_attack_features(unlearned_model, forget_loader, device, index)
+        X_eval_non, retain_latent = self._extract_attack_features(unlearned_model, retain_test_loader, device, index)
+        
+        min_test = min(len(X_eval_mem), len(X_eval_non))
+        X_test = np.vstack([X_eval_mem[:min_test], X_eval_non[:min_test]])
+        y_test = np.concatenate([np.ones(min_test), np.zeros(min_test)])
+        
+        predictions = attack_classifier.predict(X_test)
+        mia_accuracy = accuracy_score(y_test, predictions)
+        
+        clean_centroid = np.mean(retain_latent, axis=0)
+        forget_distances = np.linalg.norm(latent_features - clean_centroid, axis=1)
+        
+        return mia_accuracy, np.mean(forget_distances)
+
+    def run_logit_space_lookalike_mia(self, filtered_model, naive_retrained, forget_loader, device, target_class):
+        filtered_model.eval()
+        naive_retrained.eval()
+        
+        filtered_logits = []
+        naive_logits = []
+        
+        with torch.no_grad():
+            for data, _ in forget_loader:
+                data = data.to(device)
+                
+                if filtered_model.__class__.__name__ == "WF_Module":
+                    gate_signals = torch.full((data.size(0),), target_class, dtype=torch.long, device=data.device)
+                    out_filtered = filtered_model(data, target_class_indices=gate_signals)
+                else:
+                    out_filtered = filtered_model(data)
+                
+                out_naive = naive_retrained(data)
+                
+                filtered_logits.append(out_filtered)
+                naive_logits.append(out_naive)
+                
+        # Concatenate everything
+        filtered = torch.cat(filtered_logits, dim=0).cpu().numpy()
+        naive = torch.cat(naive_logits, dim=0).cpu().numpy()
+        
+        # Z-Score Normalisation
+        filtered = (filtered - np.mean(filtered, axis=-1, keepdims=True)) / (np.std(filtered, axis=-1, keepdims=True) + 1e-8)
+        naive = (naive - np.mean(naive, axis=-1, keepdims=True)) / (np.std(naive, axis=-1, keepdims=True) + 1e-8)
+        
+        # shuffle indices
+        num_images = len(filtered)
+        image_indices = np.arange(num_images)
+        np.random.shuffle(image_indices)
+
+        # split to train and test set
+        split = int(num_images * 0.7)
+        train_img_idx, test_img_idx = image_indices[:split], image_indices[split:]
+
+        # create a balanced test and train set
+        data_train = np.vstack([filtered[train_img_idx], naive[train_img_idx]])
+        data_test = np.vstack([filtered[test_img_idx], naive[test_img_idx]])
+
+        # labels for attcker (1 from unlearned and 0s to retrained)
+        # we do this because every output retrained gives us is a result of unseen
+        # but unlearned has seen these data.
+        label_train = np.concatenate([np.ones(len(train_img_idx)), np.zeros(len(train_img_idx))])
+        # test set
+        label_test = np.concatenate([np.ones(len(test_img_idx)), np.zeros(len(test_img_idx))])
+
+        
+        adversary = LogisticRegression(max_iter=1000)
+        adversary.fit(data_train, label_train)
+        # evaluate similarity of outputs
+        lookalike_accuracy = accuracy_score(label_test, adversary.predict(data_test))
+        # so that the metric is between 0 and 1.
+        return 2.0 * np.abs(lookalike_accuracy - 0.5)
+
+    def run_complete_evaluation(self, framework_name, target_class, forget_loader, retain_test_loader, unlearned_instance, base_shadow_instance, device):
+        """Orchestrates specific pipeline routing cleanly using cached constructor parameters."""
+        target_dir = os.path.join("reports", framework_name)
+        os.makedirs(target_dir, exist_ok=True)
+        current_log_file = os.path.join(target_dir, "attack_values.csv")
+        
+        if not os.path.exists(current_log_file):
+            with open(current_log_file, "w") as f:
+                f.write("target_class,parameter_mia_accuracy,latent_distance_tell,lookalike_accuracy\n")
+
+        self.register_model_hook(unlearned_instance.model)
+
+        # 1. Parameter-Space MIA and Latent Distance
+        parameter_mia_acc, latent_dist = self.run_parameter_space_mia(
+            unlearned_model=unlearned_instance.model,
+            shadow_model=base_shadow_instance.model,
+            forget_loader=forget_loader,
+            retain_test_loader=retain_test_loader,
+            device=device,
+            index=target_class
+        )
+
+        # we load a retrained model to evaluate look_alike tests
+        ghost_checkpoint_path = f"trained_models/class_{target_class}_retrained.pth"
+        
+        if os.path.exists(ghost_checkpoint_path):
+            # Safe clean wrapper boot utilizing internal instance state properties
+            ghost_model_instance = Model.create(arch=self.arch, device=device, size=self.class_size)
+            state_dict = torch.load(ghost_checkpoint_path, map_location=device, weights_only=True)
+            ghost_model_instance.model.load_state_dict(state_dict)
+            reference_model_torch = ghost_model_instance.model
+        else:
+            raise FileNotFoundError(
+                f"Retrained weights not found at: {ghost_checkpoint_path}. \nDid you forget to save models or are they saved with a different path?"
+            )
+
+        lookalike_acc = self.run_logit_space_lookalike_mia(
+            filtered_model=unlearned_instance.model, 
+            naive_retrained=reference_model_torch, 
+            forget_loader=forget_loader, 
+            device=device,
+            target_class=target_class
+        )
+
+        print(f"[{framework_name}] Class {target_class} | Parameter MIA: {parameter_mia_acc:.4f} | Latent Dist: {latent_dist:.4f} | Lookalike: {lookalike_acc:.4f}")
+        
+        with open(current_log_file, "a") as f:
+            f.write(f"{target_class},{parameter_mia_acc:.6f},{latent_dist:.6f},{lookalike_acc:.6f}\n")
+            
+        return {
+            "parameter_mia_accuracy": parameter_mia_acc,
+            "latent_distance": latent_dist,
+            "lookalike_accuracy": lookalike_acc
+        }