optimised

unlearning/WeightFiltration.py

@@ -0,0 +1,139 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader, ConcatDataset, Subset
+from unlearning.Strategy import Strategy
+import numpy as np
+from sklearn.metrics import classification_report
+from architectures.WFNet import WF_Net_Model
+
+from sets.Data import vertical_split
+
+class WeightFiltration(Strategy):
+    def __init__(self, 
+                 target_class_index: int, 
+                 arch,
+                 num_classes: int = 20, 
+                 epochs: int = 6, 
+                 lr: float = 100.0, 
+                 gamma: float = 0.01,
+                 lambda_1 = 25
+                 
+                 ):
+        super().__init__(target_class_index=target_class_index)
+        self.epochs = epochs
+        self.lr = lr
+        self.gamma = gamma 
+        self.num_classes = num_classes
+        self.wf_model = None 
+        self.lambda_1 = lambda_1
+        self.arch = arch
+    
+
+    def _optimise_filter(self, model: nn.Module, retain_loader: DataLoader, forget_loader: DataLoader, device) -> nn.Module:
+
+        # new WF_Model instance
+        wf_model = WF_Net_Model(
+            device=device, 
+            arch=self.arch,
+            size=self.num_classes, 
+            original_model=model, 
+            target_class_index=self.target_class_index
+        )
+
+        wf_net = wf_model.get()
+        optimizer = optim.SGD([wf_net.alpha], lr=self.lr)
+        # Use reduction='none' so we can manipulate individual item losses
+        criterion_none = nn.CrossEntropyLoss(reduction='none')
+        criterion_mean = nn.CrossEntropyLoss()
+
+        for epoch in range(self.epochs):
+            t_loss_r, t_loss_f = 0.0, 0.0
+            steps = 0
+
+            # forget and retain
+            for (r_inputs, r_labels), (f_inputs, f_labels) in zip(retain_loader, forget_loader):
+                r_inputs, r_labels = r_inputs.to(device), r_labels.to(device)
+                f_inputs, f_labels = f_inputs.to(device), f_labels.to(device)
+                
+                optimizer.zero_grad()
+                
+                # retain data paired with randomly selected rows of alpha to compute the retaining loss 
+                random_offset = torch.randint(0, self.num_classes - 1, size=r_labels.shape, device=device)
+                gate_signals_r = torch.where(random_offset >= r_labels, random_offset + 1, random_offset)
+                
+                outputs_r = wf_net(r_inputs, target_class_indices=gate_signals_r)
+                loss_r = criterion_mean(outputs_r, r_labels)
+
+                # Forget set is paired with corresponding labels as row selectors for alpha 
+                # and used to compute unlearning loss
+                outputs_f = wf_net(f_inputs, target_class_indices=f_labels)
+                
+                # Calculate loss for every single item in the batch at once
+                per_item_forget_loss = criterion_none(outputs_f, f_labels)
+                
+                # Use a scatter/sum approach to get class-wise losses without a Python loop
+                # Create a mask of unique classes present in this batch
+                unique_classes, inverse_indices = torch.unique(f_labels, return_inverse=True)
+                classes_in_batch = unique_classes.size(0)
+                
+                if classes_in_batch > 0:
+                    # average CE loss per class
+                    class_loss_sums = torch.zeros(classes_in_batch, device=device)
+                    class_loss_sums.scatter_add_(0, inverse_indices, per_item_forget_loss)
+                    
+                    class_counts = torch.zeros(classes_in_batch, device=device)
+                    class_counts.scatter_add_(0, inverse_indices, torch.ones_like(per_item_forget_loss))
+                    
+                    mean_class_ce_loss = class_loss_sums / class_counts
+                    
+                    # Poppi et al. suggest employing reciprocal of the forget loss 
+                    # to avoid shortcomings of negative gradient approach
+                    loss_f = torch.mean(1.0 / (mean_class_ce_loss + 1e-6))
+                else:
+                    loss_f = torch.tensor(0.0, device=device)
+
+                # Regularisation penalty
+                loss_reg = torch.sum(1.0 - torch.sigmoid(wf_net.alpha))
+                                
+                # Backpropagation
+                total_loss = loss_r + (self.lambda_1 * loss_f) + (self.gamma * loss_reg)
+                total_loss.backward()
+                optimizer.step()
+                                
+                # Keep tracking stats
+                t_loss_r += loss_r.item()
+                t_loss_f += loss_f.item()
+                steps += 1
+                            
+            print(f"  Epoch {epoch+1}/{self.epochs} | Retain Loss: {t_loss_r/steps:.4f} | Forget Loss: {t_loss_f/steps:.4f}")
+        
+        return wf_model
+    
+
+    def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
+        device = next(model.parameters()).device
+        model.eval()
+        
+        if self.wf_model is None:
+            print(">> Initializing and compiling global WF-Net matrix (Run Once for all classes)...")
+            
+            self.wf_model = self._optimise_filter(
+                model, 
+                retain_loader=retain_loader, 
+                forget_loader=forget_loader,
+                device=device
+            )
+        else:
+            print(f">> Gating matrix loaded. Switching layout to target class index: {self.target_class_index}")
+            self.wf_model.target_class_index = self.target_class_index
+            
+        return self.wf_model
+    
+    def _split_data(self, dataset):
+        return vertical_split(
+            dataset= dataset,
+            batch_size=32,
+            num_classes=self.num_classes
+        )
+