clean up

architectures/WFNet.py

@@ -7,6 +7,68 @@ 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_Module(nn.Module):
     def __init__(self, original_model: nn.Module, num_classes: int, arch_enum):
         super().__init__()
@@ -38,28 +100,23 @@ class WF_Module(nn.Module):
             case arch_enum.RESNET18 | arch_enum.RESNET34 | arch_enum.RESNET50 | arch_enum.WIDE_RESNET:
                 return model.layer4[-1]
                 
-            # --- GOOGLENET ---
             case arch_enum.GOOGLENET:
                 return model.inception5b
                 
-            # --- INCEPTION V3 ---
             case arch_enum.INCEPTION:
                 return model.Mixed_7c
                 
-            # --- DENSENET 121 ---
             case arch_enum.DENSENET121:
                 return model.features.norm5
                 
-            # --- EFFICIENTNET ---
             case arch_enum.EFFICIENTNET:
                 return model.features[-1]
                 
-            # --- SHUFFLENET ---
             case arch_enum.SHUFFLENET:
                 return model.conv5
                 
             case _:
-                # Robust Fallback Strategy
+                # Fallback Strategy
                 target = None
                 for module in model.modules():
                     if isinstance(module, nn.Conv2d):