certified

unlearning/LinearFiltration.py

@@ -1,47 +1,184 @@
-
 import torch
 import torch.nn as nn
 from .Strategy import Strategy
 from torch.utils.data import DataLoader
 
 class LinearFiltration(Strategy):
-    def __init__(self,target_class_index):
-        super().__init__(target_class_index = target_class_index)
+    def __init__(self, target_class_index):
+        super().__init__(target_class_index=target_class_index)
 
     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
             
-        with torch.no_grad():
-            W = model.fc.weight.data.clone()
-            num_classes = W.shape[0]
-            
-            A = self._calculate_filtration_matrix(num_classes, self.target_class_index, W.device)
-            sanitized_W = torch.mm(A, W)
-            model.fc.weight.copy_(sanitized_W)
-            # Filter the bias (if the layer uses one)
-            if model.fc.bias is not None:
-                b = model.fc.bias.data.clone()
-                # b is a 1D tensor of shape (num_classes), 
-                # so we use torch.mv (matrix-vector multiplication) or unsqueeze it
-                sanitized_b = torch.mv(A, b)
-                model.fc.bias.copy_(sanitized_b)
-            
-        return model
+        return self.normalise(
+            model=model, 
+            retain_loader=retain_loader, 
+            forget_loader=forget_loader,
+            device=device,
+            forget_index=self.target_class_index
+        )
+    
+    # FIX: Added staticmethod decorator
+    @staticmethod
+    def get_features(model, inputs):
+        # For ResNet, pass through everything up to the fc layer
+        x = model.conv1(inputs)
+        x = model.bn1(x)
+        x = model.relu(x)
+        x = model.maxpool(x)
+
+        x = model.layer1(x)
+        x = model.layer2(x)
+        x = model.layer3(x)
+        x = model.layer4(x)
+
+        x = model.avgpool(x)
+        x = torch.flatten(x, 1)
+        return x
 
     @staticmethod
     def _calculate_filtration_matrix(num_classes: int, forget_class: int, device: torch.device) -> torch.Tensor:
         A = torch.eye(num_classes, device=device)
         num_remaining = num_classes - 1
         
-        # The row of the forgotten class should average all other classes
         for j in range(num_classes):
             if j == forget_class:
-                # we zero the forget class
                 A[forget_class, j] = 0.0
             else:
-                # and we distribute the output to the remaining
                 A[forget_class, j] = 1.0 / num_remaining
                 
-        return A
+        return A
+    
+
+    @staticmethod
+    def _sums_and_counts(model, num_classes, retain_loader, forget_loader, device, forget_index, h_dim):
+        model.eval()
+
+        sums = torch.zeros(num_classes, h_dim, device=device)
+        counts = torch.zeros(num_classes, device=device)
+
+        # Generate values for retain
+        with torch.no_grad():
+            for inputs, targets in retain_loader:
+                inputs = inputs.to(device)
+                targets = targets.to(device)
+                # FIX: Call get_features instead of model() directly
+                outputs = LinearFiltration.get_features(model, inputs)
+
+                for j in range(num_classes):
+                    if j == forget_index:
+                        continue
+                    mask = (targets == j)
+
+                    if mask.any():
+                        sums[j] += outputs[mask].sum(dim=0)
+                        counts[j] += mask.sum()
+        
+        # Values for forget
+        with torch.no_grad():
+            for inputs, targets in forget_loader:
+                inputs = inputs.to(device)
+                targets = targets.to(device)
+                # FIX: Call get_features instead of model() directly
+                outputs = LinearFiltration.get_features(model, inputs)
+
+                mask = (targets == forget_index)
+
+                if mask.any():
+                    sums[forget_index] += outputs[mask].sum(dim=0)
+                    counts[forget_index] += mask.sum()
+
+        return sums, counts
+    
+    @staticmethod
+    def _get_means(model, num_classes, retain_loader, forget_loader, device, forget_index):
+        h_dim = model.fc.in_features
+
+        sums, counts = LinearFiltration._sums_and_counts(
+            model=model,
+            num_classes=num_classes,
+            retain_loader=retain_loader,
+            forget_loader=forget_loader,
+            device=device,
+            forget_index=forget_index,
+            h_dim=h_dim
+        )
+        A = []
+
+        for i in range(num_classes):
+            if counts[i] > 0:
+                A.append(sums[i] / counts[i])
+            else:
+                A.append(torch.zeros(h_dim, device=device))
+        
+        # CORRECT: Stack along dim=0 to make it (num_classes, h_dim)
+        return torch.stack(A, dim=0)
+    
+
+    @staticmethod
+    def _compute_z(tensor, forget_index):
+        # Now tensor has shape (num_classes, h_dim) -> tensor.shape[0] is num_classes
+        K = tensor.shape[0] 
+        
+        # pi_a0 should match the feature space dimensions (h_dim)
+        pi_a0 = torch.zeros(tensor.shape[1], device=tensor.device)
+
+        t_1 = pi_a0
+        a0 = tensor[forget_index, :] # Extracting the row vector for the forgotten class
+
+        mask_a0 = torch.ones(
+            a0.shape[0],
+            dtype=torch.bool, 
+            device=tensor.device
+        )
+        # We compute the target shift over features
+        t_2 = -(1.0 / (K - 1)) * a0[mask_a0].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
+
+
+    @staticmethod
+    def normalise(model, retain_loader, forget_loader, device, forget_index):
+        W = model.fc.weight.data.clone()
+        num_classes = W.shape[0]
+        
+        A = LinearFiltration._get_means(
+            model=model,
+            num_classes=num_classes,
+            retain_loader=retain_loader,
+            forget_loader=forget_loader,
+            device=device,
+            forget_index=forget_index
+        )
+
+        Z = LinearFiltration._compute_z(tensor=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(A[i])
+        
+        # Stack back along dim=0 to match (num_classes, h_dim)
+        B_Z = torch.stack(B_Z_rows, dim=0)
+
+        A_inv = torch.linalg.pinv(A)
+        
+        W_Z = B_Z @ A_inv @ W
+
+        model.fc.weight.copy_(W_Z)
+
+        return model