reports from optimised linear filtration

unlearning/LinearFiltration.py

@@ -2,12 +2,13 @@ import torch
 import torch.nn as nn
 from .Strategy import Strategy
 from torch.utils.data import DataLoader
-from sets.Data import get_unlearning_loaders, _combine_set
+from sets.Data import get_unlearning_loaders, _combine_set, vertical_split
 
 class LinearFiltration(Strategy):
-    def __init__(self, target_class_index):
+    def __init__(self, target_class_index, num_classes = 20):
         super().__init__(target_class_index=target_class_index)
         self.A = None
+        self.num_classes = num_classes
 
     def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
         model.eval()
@@ -51,13 +52,13 @@ class LinearFiltration(Strategy):
     
   
 
-    def _compute_A(self, model, num_classes, loader, device):
+    def _compute_A(self, model, loader, device):
         model.eval()
         
         
         # Initialize tracking tensors
-        sums = torch.zeros(num_classes, num_classes, device=device)
-        counts = torch.zeros(num_classes, device=device)
+        sums = torch.zeros(self.num_classes, self.num_classes, device=device)
+        counts = torch.zeros(self.num_classes, device=device)
 
         with torch.no_grad():
             for inputs, targets in loader:
@@ -67,7 +68,7 @@ class LinearFiltration(Strategy):
                 outputs = model(inputs) 
                 
                 # One-hot encode targets to act as a routing mask
-                one_hot = torch.nn.functional.one_hot(targets, num_classes=num_classes).float()
+                one_hot = torch.nn.functional.one_hot(targets, num_classes=self.num_classes).float()
                 
                 # add 
                 sums += torch.t(one_hot) @ outputs
@@ -77,7 +78,6 @@ class LinearFiltration(Strategy):
 
         # means
         counts_safe = counts.unsqueeze(1)
-        print(f"COUNTS IS >>>>>>>>> {counts_safe}")
         self.A = torch.where(
             counts_safe > 0, 
             sums / counts_safe, 
@@ -117,10 +117,10 @@ class LinearFiltration(Strategy):
     def normalise(self, model, retain_loader, forget_loader, device, forget_index):
         clf = self._get_classifier(model)
         W = clf.weight.data.clone()
-        num_classes = W.shape[0]
+        #num_classes = W.shape[0]
 
         # we combine the data so we can calculate the mean of prdictions
-        full_loader = _combine_set(retain_loader, forget_loader)
+        #full_loader = _combine_set(retain_loader, forget_loader)
         # 8
         # Computing A is the most resource intensive part of this algorithm
         # and to optimise the process, we computr it only once and re-use it
@@ -128,8 +128,8 @@ class LinearFiltration(Strategy):
         if self.A is None:
             self._compute_A(
                 model = model,
-                num_classes = num_classes,
-                loader = full_loader,
+                #num_classes = num_classes,
+                loader = forget_loader,
                 device = device
             )
 
@@ -137,7 +137,7 @@ class LinearFiltration(Strategy):
         Z = self._compute_z(tensor=self.A, forget_index=forget_index)
         B_Z_rows = []
 
-        for i in range(num_classes):
+        for i in range(self.num_classes):
             if i == forget_index:
                 B_Z_rows.append(Z)
             else:
@@ -161,8 +161,14 @@ class LinearFiltration(Strategy):
     
     # overriden function
     def _split_data(self, dataset):
-        return get_unlearning_loaders(
+        '''return get_unlearning_loaders(
             dataset=dataset,
             forget_class_idx=self.target_class_index,
             batch_size = 32
+        )'''
+        return vertical_split(
+            dataset= dataset,
+            batch_size=32,
+            num_classes=self.num_classes,
+            ratio=0.1
         )