unlearning done

unlearning/wf/WF_Net.py

@@ -35,52 +35,52 @@ class WF_Net(nn.Module):
         #self.alpha = nn.Parameter(torch.ones(num_classes, out_channels) * 1.5)
         self.alpha = nn.Parameter(torch.ones(num_classes, out_channels))
 
-    def forward(self, x: torch.Tensor, target_unlearn_class: int) -> torch.Tensor:
-        """
-        Implements Algorithm 1: General forward step of a WF model
-        Inputs: 
-          x: Input tensor (Xin)
-          target_unlearn_class: The class index we are actively filtering out (Yunl)
-        """
+    def forward(self, x: torch.Tensor, target_class_indices: torch.Tensor) -> torch.Tensor:
         # 1. Run through early sequence of layers undisturbed
         x = self.maxpool(self.relu(self.bn1(self.conv1(x))))
         x = self.layer1(x)
         x = self.layer2(x)
         x = self.layer3(x)
-        
-        # Run layer4 block 0 and block 1 conv1 normally
+                
+        # Run layer4 block 0 normally
         x = self.layer4[0](x)
-
-        identity = x
-
-
-
+        
+        # -------------------------------------------------------------
+        # HERE IT IS: Save the structural skip connection (identity) 
+        # BEFORE modifying features via block 1's convolutions
+        # -------------------------------------------------------------
+        identity = x 
+                
+        # Now enter layer4 block 1
         x = self.layer4[1].conv1(x)
         x = self.layer4[1].bn1(x)
         x = self.layer4[1].relu(x)
+                
+        # [Your Step 1 Masking Math happens right here...]
+        batch_alpha = self.alpha[target_class_indices]
+        mask = torch.sigmoid(batch_alpha).view(x.size(0), -1, 1, 1)
         
-        # 2. CORE WF-NET MATH: w_hat_l <- alpha_l[Yunl] ⊙ w_l
-        # Extract 1D vector for target class and reshape to (out_channels, 1, 1, 1) for 4D convolution broadcasting
-        mask = torch.sigmoid(self.alpha[target_unlearn_class]).view(-1, 1, 1, 1)
-        w_hat = self.original_w * mask
-        
-        # 3. Pass gated weights straight to functional forward pass: l(Xi, w_hat_l)
+        # Run the functional convolution
         x = F.conv2d(
             x, 
-            weight=w_hat, 
+            weight=self.original_w, 
             bias=self.target_conv.bias, 
             stride=self.target_conv.stride, 
             padding=self.target_conv.padding
         )
+        
+        # Apply your WF-Net channel mask
+        x = x * mask 
         x = self.layer4[1].bn2(x)
-        
-        # Handle residual shortcut skip connection manually since we opened up block 1
-        # In ResNet-18 layer4, block 1 has no downsample shortcut layer; it's a direct identity add
+                
+        # -------------------------------------------------------------
+        # HERE IT IS USED: Add the pristine identity back to the gated output
+        # -------------------------------------------------------------
         x = self.layer4[1].relu(x + identity) 
-        
-        # 4. Final Classification Head Sequence
+                
+        # Final Classification Head Sequence
         x = self.avgpool(x)
         x = torch.flatten(x, 1)
         y_out = self.fc(x)
-        
+                
         return y_out