facebook's implementation

unlearning/WeightFiltration.py

@@ -1,6 +1,5 @@
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 import torch.optim as optim
 from torch.utils.data import DataLoader
 from unlearning.Strategy import Strategy
@@ -10,10 +9,9 @@ class WeightFiltration(Strategy):
     Implements Poppi et al.'s Weight Filtering framework for linear layers.
     Uses a standard functional hook to guarantee native PyTorch autograd tracking.
     """
-    def __init__(self, num_classes: int, target_class_idx: int, epochs: int = 10, lr: float = 0.2, gamma: float = 10.0):
-        super().__init__()
+    def __init__(self, target_class_index,num_classes: int, epochs: int = 10, lr: float = 0.2, gamma: float = 10.0):
+        super().__init__(target_class_index = target_class_index)
         self.num_classes = num_classes
-        self.target_class_idx = target_class_idx
         self.epochs = epochs
         self.lr = lr
         self.gamma = gamma 
@@ -52,13 +50,13 @@ class WeightFiltration(Strategy):
         
         # Transfer the channel suppression permanently into model.fc
         with torch.no_grad():
-            mask = torch.sigmoid(self.alpha[self.target_class_idx]) # Shape: (num_features,)
+            mask = torch.sigmoid(self.alpha[self.target_class_index]) # Shape: (num_features,)
             
             # Suppress the channels ONLY for the target class row in fc
-            fc_layer.weight[self.target_class_idx].copy_(
-                fc_layer.weight[self.target_class_idx] * mask
+            fc_layer.weight[self.target_class_index].copy_(
+                fc_layer.weight[self.target_class_index] * mask
             )
-            print(f">> Baked deep channel filter into Class {self.target_class_idx} weights.")
+            print(f">> Baked deep channel filter into Class {self.target_class_index} weights.")
             
         return model
 
@@ -72,7 +70,7 @@ class WeightFiltration(Strategy):
             # self.alpha shape: (num_classes, channels) -> e.g., (20, 2048)
             
             # Extract 1D mask for the target class: (channels,)
-            mask = torch.sigmoid(self.alpha[self.target_class_idx])
+            mask = torch.sigmoid(self.alpha[self.target_class_index])
             
             # Reshape mask to (1, channels, 1, 1) so it broadcasts over batch, height, and width
             mask = mask.view(1, -1, 1, 1)
@@ -87,7 +85,7 @@ class WeightFiltration(Strategy):
         optimizer = optim.Adam([self.alpha], lr=self.lr)
         criterion = nn.CrossEntropyLoss()
 
-        print(f"[{self.__class__.__name__}] Unlearning Class {self.target_class_idx} with gamma={self.gamma}...")
+        print(f"[{self.__class__.__name__}] Unlearning Class {self.target_class_index} with gamma={self.gamma}...")
 
         # To optimise this loop we will watch improvements after each optimisation
         temp_forget_loss = None