cleaned up for submission

architectures/WFNet.py

@@ -6,69 +6,6 @@ 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_Module(nn.Module):
     def __init__(self, original_model: nn.Module, num_classes: int, arch_enum):
         super().__init__()