unlearning done

unlearning/WeightFiltration.py

@@ -1,126 +1,135 @@
 import torch
 import torch.nn as nn
 import torch.optim as optim
-from torch.utils.data import DataLoader
+from torch.utils.data import DataLoader, ConcatDataset, Subset
 from unlearning.Strategy import Strategy
-from .wf.WF_Net import WF_Net
+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):
-    """
-    Verbatim implementation of Poppi et al.'s WF-Net framework.
-    Directly filters the convolutional weights of a target layer using a learnable
-    channel mask, optimizing it via weight-space regularization.
-    """
-    def __init__(self, target_class_index: int, epochs: int = 10, lr: float = 0.2, gamma: float = 10.0):
+    def __init__(self, 
+                 target_class_index: int, 
+                 num_classes: int = 20, 
+                 epochs: int = 6, 
+                 lr: float = 100.0, 
+                 gamma: float = 0.01,
+                 ):
         super().__init__(target_class_index=target_class_index)
         self.epochs = epochs
         self.lr = lr
         self.gamma = gamma 
-        #self.alpha = None
+        self.num_classes = num_classes
+        self.wf_model = None 
+        self.lambda_1 = 25
+    
+    
+    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, 
+            size=self.num_classes, 
+            original_model=model, 
+            target_class_index=self.target_class_index
+        )
 
-
-
-
-    def _optimise_filter(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader, device):
-        # 1. Initialize the wrapper with your pre-trained model
-        num_classes = model.fc.out_features
-        wf_model = WF_Net(original_model=model, num_classes=num_classes).to(device)
-        
-        # 2. ONLY optimize alpha (everything else is frozen inside the wrapper)
-        optimizer = optim.Adam([wf_model.alpha], lr=self.lr)
+        # a WF_net module to be trained (unlearned) to generate alpha
+        wf_net = wf_model.get()
+        optimizer = optim.SGD([wf_net.alpha], lr=self.lr)
         criterion = nn.CrossEntropyLoss()
 
         for epoch in range(self.epochs):
-            forget_iter = iter(forget_loader)
             t_loss_r, t_loss_f = 0.0, 0.0
             steps = 0
             
-            for r_inputs, r_labels in retain_loader:
+            # 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)
-                
-                # Pull the matching forget batch input
-                try:
-                    f_inputs, _ = next(forget_iter)
-                except StopIteration:
-                    forget_iter = iter(forget_loader)
-                    f_inputs, _ = next(forget_iter)
-                f_inputs = f_inputs.to(device)
+                f_inputs, f_labels = f_inputs.to(device), f_labels.to(device)
                 
                 optimizer.zero_grad()
                 
-                # --- APPLY ALGORITHM 1 FORWARD PASS TO BOTH INPUTS ---
-                # Pass the input batch AND the target unlearn class index
-                outputs_r = wf_model(r_inputs, target_unlearn_class=self.target_class_index)
-                outputs_f = wf_model(f_inputs, target_unlearn_class=self.target_class_index)
+                # retain data paired with randomly selected rows of alpha to compute the retaining loss 
+                random_rows = []
+                for label in r_labels:
+                    allowed = [i for i in range(self.num_classes) if i != label.item()]
+                    random_rows.append(np.random.choice(allowed))
                 
-                # Compute Losses using Poppi et al.'s temperature scaled entropy
+                gate_signals_r = torch.tensor(random_rows, dtype=torch.long, device=device)
+                outputs_r = wf_net(r_inputs, target_class_indices=gate_signals_r)
                 loss_r = criterion(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)
                 
-                temperature = 3.0
-                logits_f_scaled = outputs_f / temperature
+                loss_f = 0.0
+                classes_in_batch = 0 
                 
-                # Compute uniform target entropy per-sample, then average over the batch
-                log_probs_f = torch.log_softmax(logits_f_scaled, dim=-1)
-                uniform_target = torch.ones_like(logits_f_scaled) / num_classes
-                loss_f = -torch.sum(uniform_target * log_probs_f, dim=-1).mean()
+                # every image of class c will unlearn over the same row of alpha_l (poppi et al page 5)
+                for c in range(self.num_classes):
+                    class_mask = (f_labels == c)
+                    if not class_mask.any():
+                        continue
+                        
+                    labels_c = f_labels[class_mask]
+                    
+                    # Slice the existing outputs instead of recalculating a forward pass
+                    outputs_f_c = outputs_f[class_mask] 
+                    
+                    loss_f_ce = criterion(outputs_f_c, labels_c)
+                    
+                    # Poppi et al. suggest employing reciprocal of the forget loss 
+                    # to avoid shortcomings of negative gradient approach
+                    loss_f += 1.0 / (loss_f_ce + 1e-6)
+                    classes_in_batch += 1
+                
+                # Average forget loss by number of distinct classes seen in this batch
+                if classes_in_batch > 0:
+                    loss_f = loss_f / classes_in_batch
+
+                # Regilarisation penality
+                loss_reg = torch.sum(1.0 - torch.sigmoid(wf_net.alpha))
                                 
-                total_loss = loss_r + (self.gamma * loss_f)
+                # back propagation
+                total_loss = loss_r + (self.lambda_1 * loss_f) + (self.gamma * loss_reg)
                 total_loss.backward()
                 optimizer.step()
                                 
                 t_loss_r += loss_r.item()
-                t_loss_f += loss_f.item()
+                t_loss_f += loss_f.item() if classes_in_batch > 0 else 0.0
                 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() 
+        model.eval()
         
-        # In WF-Net, the mask targets the last major convolutional block
-        # For ResNet-18, that is the final conv layer in layer4 block 1
-        if hasattr(model, 'layer4') and len(model.layer4) > 1:
-            target_conv = model.layer4[1].conv2
+        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:
-            raise AttributeError("Model architecture does not match expected ResNet-18 structure.")
+            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
             
-        # Store a pristine, non-grad copy of the original trained weights
-        # Shape of conv2.weight: (out_channels, in_channels, kernel_size, kernel_size) -> e.g., (512, 512, 3, 3)
-        original_weights = target_conv.weight.data.clone().detach()
-        out_channels = original_weights.shape[0]
-
-        # Initialize alpha gate vector matching Poppi et al.'s initialization range
-        # Shape: (out_channels,) -> acting directly as a filter-level gate
-        #self.alpha = nn.Parameter(torch.ones(out_channels, device=device) * 1.5)
-
-        # Freeze the global model graph; only optimize our filter parameter mask
-        for p in model.parameters():
-            p.requires_grad = False
-        #self.alpha.requires_grad = True
-
-        wf_model = self._optimise_filter(
-            model,
-            forget_loader=forget_loader,
-            retain_loader=retain_loader,
-            device=device,
+        return self.wf_model
+    
+    def _split_data(self, dataset):
+        return vertical_split(
+            dataset= dataset,
+            batch_size=32,
+            num_classes=self.num_classes
         )
-
-        # --- PERMANENT BAKING STEP ---
-        with torch.no_grad():
-            # Grab the alpha mask vector for the forgotten class and cast to 4D tensor shape
-            final_mask = torch.sigmoid(wf_model.alpha[self.target_class_index]).view(-1, 1, 1, 1)
-            
-            # Apply filter masking permanently back onto the base layer
-            target_conv.weight.copy_(original_weights * final_mask)
-                        
-            # Unfreeze architecture parameters for evaluations downstream
-            for p in model.parameters():
-                p.requires_grad = True
-                        
-            print(f">> Permanently altered {out_channels} convolutional filters in layer4 via WF-Net.")
-        return model
+