WF net added

Tune_new.py

@@ -20,7 +20,7 @@ BATCH_SIZE = 16
 SAMPLE_SIZE = 30
 TRAINING_SAMPLE = 27
 RESOLUTION = 224
-ARCH = Architecture.RESNET18
+ARCH = Architecture.GOOGLENET
 
 
 # Data preparation and model setup
@@ -230,8 +230,8 @@ if __name__ == "__main__":
 
         strategies = [
             certified_unlearning,
-            weight_filtration,
-            linear_filtration
+            #weight_filtration,
+            #linear_filtration
         ]
         # Unlearning Iteration
         for i in range(0, CLASS_SIZE):

architectures/WFNet.py

@@ -0,0 +1,126 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torch.utils.data import DataLoader
+import numpy as np
+from sklearn.metrics import classification_report
+from architectures.Model import Model
+
+class WF_Module(nn.Module):
+    """
+    Pure PyTorch Neural Network module graph.
+    Keeps parameter registration and autograd tracking separate from 
+    the framework's high-level Model abstractions to prevent recursion collisions.
+    """
+    def __init__(self, original_model: nn.Module, num_classes: int):
+        super().__init__()
+    
+        self.original_model = original_model
+        
+        # Target layer for weight filtering (layer4 block 1 conv2 or conv3 depending on arch)
+        last_layer = original_model.layer4[1]
+
+        # Some versions are limited to 2 convolutional layers
+        if hasattr(last_layer, "conv3"):
+            self.target_conv = last_layer.conv3
+        else:
+            self.target_conv = last_layer.conv2
+
+        # Completely freeze the original ResNet parameters
+        for param in self.parameters():
+            param.requires_grad = False
+            
+        # Initialize the alpha parameter matrix (Rows = Classes, Cols = Channels)
+        out_channels = self.target_conv.weight.shape[0]
+        self.alpha = nn.Parameter(torch.full((num_classes, out_channels), 3.0))
+
+    '''
+    Poppi et_al's Single-shot multiclass unlearning.
+    This calculation happens only once to generate the mask. once the mask is generated,
+    Unlearning and remembering becomes a matter of switching gates on and off.
+    '''
+    def forward(self, x: torch.Tensor, target_class_indices: torch.Tensor) -> torch.Tensor:
+        # we linearly loop through layers 1 to 4[block 1] (for ResNet) 
+        # for i in M_{|L|} do  l <- l[i]
+        x = self.original_model.maxpool(self.original_model.relu(self.original_model.bn1(self.original_model.conv1(x))))
+        x = self.original_model.layer1(x)
+        x = self.original_model.layer2(x)
+        x = self.original_model.layer3(x)
+        x = self.original_model.layer4[0](x)
+        
+        # The second block execute its internal transformations natively
+        # This handles conv1->conv2 (ResNet18) or conv1->conv2->conv3 (ResNet50) automatically!
+        # Xi+1 <- l(Xi, ˆwl)
+        x = self.original_model.layer4[1](x)
+        
+        # Apply mask dynamically to the completed block feature map
+        # wl <- αl[Yunl] ⊙ ˆwl
+        batch_alpha = self.alpha[target_class_indices]
+        mask = torch.sigmoid(batch_alpha).view(x.size(0), -1, 1, 1)
+        x = x * mask
+        
+        # Remaining standard head steps
+        x = self.original_model.avgpool(x)
+        x = torch.flatten(x, 1)
+        # so here we are returning the output logits
+        # the result of classification is then 
+        # argmax(x)
+        return self.original_model.fc(x)
+    
+
+
+class WF_Net_Model(Model):
+    def __init__(self, device, size, original_model: nn.Module, target_class_index: int):
+        self.device = device
+        self.size = size
+        self.wf_module = WF_Module(
+            original_model = original_model, 
+            num_classes = size
+        ).to(self.device)
+        
+        # this index indicates which row of the mask should be active (gate closed).
+        self.target_class_index = target_class_index
+        self.model = self.wf_module
+
+    def get(self):
+        return self.wf_module
+    
+    '''
+    We override the evaluate method from the base class, 
+    because how we evaluate is different here from that of a normal torch nn.Module object
+
+    '''
+    def evaluate(self, loader, mode="eval"):
+
+        self.wf_module.eval()
+        all_preds, all_labels = [], []
+        print(f"\nEvaluating Domain: [{mode}]...")
+        
+        with torch.no_grad():
+            for inputs, labels in loader:
+                inputs, labels = inputs.to(self.device), labels.to(self.device)
+                
+                # we apply the filter
+                gate_signals = torch.full((inputs.size(0),), self.target_class_index, dtype=torch.long, device=self.device)
+                
+                # pass prediction through the filter
+                outputs = self.wf_module(inputs, target_class_indices=gate_signals)
+
+                # return argmax(x)    
+                _, predicted = torch.max(outputs, 1)
+                all_preds.extend(predicted.cpu().numpy())
+                all_labels.extend(labels.cpu().numpy())
+
+        classes = sorted(list(set(all_labels)))
+        accuracy = 100 * (np.array(all_preds) == np.array(all_labels)).sum() / len(all_labels)
+        
+        print(f"Test Accuracy: {accuracy:.2f}%")
+        print(classification_report(all_labels, all_preds, labels=classes, zero_division=0))
+        report = classification_report(all_labels, all_preds, labels=classes, output_dict=True, zero_division=0)
+        
+        return accuracy, report
+    
+    def eval(self):
+        """Safely intercept any fallback base class calls targeting .eval()"""
+        self.wf_module.eval()