From 0a7a2e1da5e089755dd046889f7c7505fe5863f5 Mon Sep 17 00:00:00 2001 From: Tinsae Date: Sun, 28 Jun 2026 01:48:49 +0200 Subject: [PATCH] WF net added --- Tune_new.py | 6 +- architectures/WFNet.py | 126 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 129 insertions(+), 3 deletions(-) create mode 100644 architectures/WFNet.py diff --git a/Tune_new.py b/Tune_new.py index 0145b9a..fdf5688 100644 --- a/Tune_new.py +++ b/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): diff --git a/architectures/WFNet.py b/architectures/WFNet.py new file mode 100644 index 0000000..96ff3a6 --- /dev/null +++ b/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()