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_Module(nn.Module): def __init__(self, original_model: nn.Module, num_classes: int, arch_enum): super().__init__() # If your model classes contain the raw inner torch model under an attribute, # extract it. Otherwise, use it directly. self.original_model = getattr(original_model, "model", original_model) # Freeze the original model parameters completely for param in self.original_model.parameters(): param.requires_grad = False # Target layer discovery using your clean Enum contract self.target_layer = self._deduce_target_layer(self.original_model, arch_enum) # Derive channel dimensions dynamically from the deduced layer out_channels = self._extract_channels(self.target_layer, self.original_model) # Initialize alpha parameter matrix (Rows = Classes, Cols = Channels) self.alpha = nn.Parameter(torch.full((num_classes, out_channels), 3.0)) self._current_target_indices = None def _deduce_target_layer(self, model: nn.Module, arch_enum) -> nn.Module: """ Scans the architecture topology to target the final deep feature extraction block right before global pooling/classification using strict Enum configurations. """ match arch_enum: # --- RESNET FAMILY --- 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 target = None for module in model.modules(): if isinstance(module, nn.Conv2d): target = module if target is not None: return target raise RuntimeError(f"Could not locate filtration anchor for Enum target: {arch_enum}") def _extract_channels(self, target_layer: nn.Module, model: nn.Module) -> int: """Helper to determine channel depth across varied layers types.""" if hasattr(target_layer, "out_channels"): return target_layer.out_channels if hasattr(target_layer, "num_features"): return target_layer.num_features if hasattr(target_layer, "weight"): return target_layer.weight.shape[0] # Classifier fallback mapping if hasattr(model, "fc"): return model.fc.in_features if hasattr(model, "classifier"): if isinstance(model.classifier, nn.Linear): return model.classifier.in_features if isinstance(model.classifier, nn.Sequential): return model.classifier[0].in_features return 512 def _filtration_hook(self, module: nn.Module, hook_input: tuple, hook_output: torch.Tensor) -> torch.Tensor: if self._current_target_indices is None: return hook_output batch_alpha = self.alpha[self._current_target_indices] if len(hook_output.shape) == 4: mask = torch.sigmoid(batch_alpha).view(hook_output.size(0), -1, 1, 1) else: mask = torch.sigmoid(batch_alpha).view(hook_output.size(0), -1) return hook_output * mask def forward(self, x: torch.Tensor, target_class_indices: torch.Tensor) -> torch.Tensor: self._current_target_indices = target_class_indices hook_handle = self.target_layer.register_forward_hook(self._filtration_hook) try: logits = self.original_model(x) finally: hook_handle.remove() self._current_target_indices = None return logits class WF_Net_Model(Model): def __init__(self, device, size, original_model: nn.Module, target_class_index: int, arch): self.device = device self.size = size self.wf_module = WF_Module( arch_enum=arch, 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()