optimised

unlearning/LinearFiltration.py

@@ -0,0 +1,168 @@
+import torch
+import torch.nn as nn
+from .Strategy import Strategy
+from torch.utils.data import DataLoader
+from sets.Data import get_unlearning_loaders, _combine_set
+
+class LinearFiltration(Strategy):
+    def __init__(self, target_class_index):
+        super().__init__(target_class_index=target_class_index)
+        self.A = None
+
+    def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
+        model.eval()
+        # Freeze internal params
+        for param in model.parameters():
+            param.requires_grad = False
+
+        device = next(model.parameters()).device
+            
+        return self.normalise(
+            model=model, 
+            retain_loader=retain_loader, 
+            forget_loader=forget_loader,
+            device=device,
+            forget_index=self.target_class_index
+        )
+    
+
+    def _get_classifier(self, model: nn.Module) -> nn.Linear:
+
+        inner_model = getattr(model, "model", model)
+        
+        # looking for standard naming conventions in named modules
+        for name, module in inner_model.named_modules():
+            # Check if it's our target linear layer
+            if (name == "fc" or name == "classifier") and isinstance(module, nn.Linear):
+                return module
+                
+            # Handle models (like EfficientNet) where the classifier is a Sequential block
+            if name == "classifier" and isinstance(module, nn.Sequential):
+                for sub_module in reversed(list(module.children())):
+                    if isinstance(sub_module, nn.Linear):
+                        return sub_module
+
+        # scan backwards for the last Linear layer
+        for module in reversed(list(inner_model.modules())):
+            if isinstance(module, nn.Linear):
+                return module
+                
+        raise RuntimeError(f"Could not locate a linear classification head for {model.__class__.__name__}")
+    
+  
+
+    def _compute_A(self, model, num_classes, loader, device):
+        model.eval()
+        
+        
+        # Initialize tracking tensors
+        sums = torch.zeros(num_classes, num_classes, device=device)
+        counts = torch.zeros(num_classes, device=device)
+
+        with torch.no_grad():
+            for inputs, targets in loader:
+                inputs, targets = inputs.to(device), targets.to(device)
+                
+                # the logit predictions
+                outputs = model(inputs) 
+                
+                # One-hot encode targets to act as a routing mask
+                one_hot = torch.nn.functional.one_hot(targets, num_classes=num_classes).float()
+                
+                # add 
+                sums += torch.t(one_hot) @ outputs
+                
+                # Sum columns of one-hot to get counts per class in this batch
+                counts += one_hot.sum(dim=0)
+
+        # means
+        counts_safe = counts.unsqueeze(1)
+        print(f"COUNTS IS >>>>>>>>> {counts_safe}")
+        self.A = torch.where(
+            counts_safe > 0, 
+            sums / counts_safe, 
+            torch.zeros_like(sums)
+        )
+
+    # 9
+    def _compute_z(self, tensor, forget_index):
+
+        K = tensor.shape[0] 
+        
+        
+        pi_a_f = torch.zeros(tensor.shape[1], device=tensor.device)
+
+        t_1 = pi_a_f
+        # row vector for the forgotten class
+        a_f = tensor[forget_index, :] 
+
+        mask_a_f = torch.ones(
+            a_f.shape[0],
+            dtype=torch.bool, 
+            device=tensor.device
+        )
+        # We compute the target shift over features
+        t_2 = -(1.0 / (K - 1)) * a_f[mask_a_f].sum()
+
+        mask_rows = torch.ones(K, dtype=torch.bool, device=tensor.device)
+        mask_rows[forget_index] = False
+
+        r_A = tensor[mask_rows, :]
+        t_3 = (1.0 / ((K - 1)) ** 2) * r_A.sum()
+
+        return t_1 + t_2 + t_3
+
+
+    # Normalisation filtration
+    def normalise(self, model, retain_loader, forget_loader, device, forget_index):
+        clf = self._get_classifier(model)
+        W = clf.weight.data.clone()
+        num_classes = W.shape[0]
+
+        # we combine the data so we can calculate the mean of prdictions
+        full_loader = _combine_set(retain_loader, forget_loader)
+        # 8
+        # Computing A is the most resource intensive part of this algorithm
+        # and to optimise the process, we computr it only once and re-use it
+        # because mean of all prdictions is the same for all
+        if self.A is None:
+            self._compute_A(
+                model = model,
+                num_classes = num_classes,
+                loader = full_loader,
+                device = device
+            )
+
+        # 9
+        Z = self._compute_z(tensor=self.A, forget_index=forget_index)
+        B_Z_rows = []
+
+        for i in range(num_classes):
+            if i == forget_index:
+                B_Z_rows.append(Z)
+            else:
+                # Retained classes maintain their original ideal feature directions
+                B_Z_rows.append(self.A[i])
+        
+        # 10
+        # Stack back along dim=0 to match (num_classes, h_dim)
+        # to get mean
+        B_Z = torch.stack(B_Z_rows, dim=0)
+
+        A_inv = torch.linalg.pinv(self.A)
+        # 11 
+        W_Z = B_Z @ A_inv @ W
+
+        # 12
+        clf = self._get_classifier(model)
+        clf.weight.copy_(W_Z)
+
+        return model
+    
+    # overriden function
+    def _split_data(self, dataset):
+        return get_unlearning_loaders(
+            dataset=dataset,
+            forget_class_idx=self.target_class_index,
+            batch_size = 32
+        )