optimised

unlearning/WF.py

@@ -0,0 +1,107 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader
+from unlearning.Strategy import Strategy
+from .wf.WF_Net import WF_Net
+
+class WeightF(Strategy):
+    """
+    Verbatim implementation of Poppi et al.'s WF-Net framework modified
+    for static, single-class unlearning extraction.
+    """
+    def __init__(self, target_class_index: int, epochs: int = 10, lr: float = 0.2, gamma: float = 10.0):
+        super().__init__(target_class_index=target_class_index)
+        self.epochs = epochs
+        self.lr = lr
+        self.gamma = gamma 
+
+    def _optimise_filter(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader, device):
+        num_classes = model.fc.out_features
+        wf_model = WF_Net(original_model=model, num_classes=num_classes).to(device)
+                
+        # Optimize only the specific alpha masks
+        optimizer = optim.Adam([wf_model.alpha], lr=self.lr)
+        criterion = nn.CrossEntropyLoss() # Default reduction is 'mean'
+        
+        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:
+                r_inputs, r_labels = r_inputs.to(device), r_labels.to(device)
+                                
+                try:
+                    f_inputs, _ = next(forget_iter)
+                except StopIteration:
+                    forget_iter = iter(forget_loader)
+                    f_inputs, _ = next(forget_iter)
+                f_inputs = f_inputs.to(device)
+                                
+                optimizer.zero_grad()
+                                
+                # Forward Pass
+                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 Loss (Mean over batch)
+                loss_r = criterion(outputs_r, r_labels)
+                                
+                # Forget Loss (Corrected to Mean over batch)
+                temperature = 1.0
+                logits_f_scaled = outputs_f / temperature
+                
+                # 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()
+                                
+                total_loss = loss_r + (self.gamma * loss_f)
+                total_loss.backward()
+                optimizer.step()
+                                
+                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 hasattr(model, 'layer4') and len(model.layer4) > 1:
+            target_conv = model.layer4[1].conv2
+        else:
+            raise AttributeError("Model architecture does not match expected ResNet-18 structure.")
+                    
+        original_weights = target_conv.weight.data.clone().detach()
+        out_channels = original_weights.shape[0]
+        
+        # Freeze global network layers
+        for p in model.parameters():
+            p.requires_grad = False
+            
+        wf_model = self._optimise_filter(
+            model,
+            forget_loader=forget_loader,
+            retain_loader=retain_loader,
+            device=device,
+        )
+        
+        # --- 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