strategies tested

unlearning/LinearFiltration.py

@@ -2,13 +2,14 @@
 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_idx: int):
-        super().__init__()  # Automatically configures 'NormalizingLinearFiltration_metrics.txt'
+        super().__init__()
         self.target_class_idx = target_class_idx
 
-    def _run(self, model: nn.Module) -> nn.Module:
+    def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
         model.eval()
         for param in model.parameters():
             param.requires_grad = False
@@ -20,29 +21,28 @@ class LinearFiltration(Strategy):
             A = self._calculate_filtration_matrix(num_classes, self.target_class_idx, 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
 
-    '''@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_classes = num_classes - 1
-        for j in range(num_classes):
-            if j == forget_class:
-                A[forget_class, j] = 0.0
-            else:
-                A[forget_class, j] = 1.0 / num_remaining_classes
-        return A'''
     @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 the inputs of all other classes
+        # 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