From 31f461342e6e70ba8d6475e3d4028a1e70486fb6 Mon Sep 17 00:00:00 2001 From: Tinsae Date: Wed, 8 Jul 2026 20:36:49 +0200 Subject: [PATCH] clean up --- architectures/WFNet.py | 69 ++++++++++++++-- eval/UnlearningAttack.py | 144 ++------------------------------- unlearning/LinearFiltration.py | 4 - 3 files changed, 70 insertions(+), 147 deletions(-) diff --git a/architectures/WFNet.py b/architectures/WFNet.py index bd06a19..4d0db37 100644 --- a/architectures/WFNet.py +++ b/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): diff --git a/eval/UnlearningAttack.py b/eval/UnlearningAttack.py index 5497bf3..d2d8ca5 100644 --- a/eval/UnlearningAttack.py +++ b/eval/UnlearningAttack.py @@ -17,10 +17,6 @@ class UnlearningAttack: """ self.arch = arch self.class_size = class_size - - #self.hook = None - #self.model = None - #self._hook_features = [] self.criterion = nn.CrossEntropyLoss(reduction='none') self.collecting = False @@ -45,109 +41,14 @@ class UnlearningAttack: return np.mean(jensenshannon(np.array(probs1), np.array(probs2), axis=1)) def calculate_a_dist(self, latent1, latent2): - """Calculates formal A-Distance: 2 * (1 - 2 * epsilon).""" - '''combined = np.vstack([latent1, latent2]) - mean = np.mean(combined, axis=0) - std = np.std(combined, axis=0) + 1e-8 - l1 = (latent1 - mean) / std - l2 = (latent2 - mean) / std - - # 2. Use the same balanced split and regularization (C=0.01) - # as the look-alike method to ensure stability. - X = np.vstack([l1, l2]) - y = np.concatenate([np.ones(len(l1)), np.zeros(len(l2))]) - - # Shuffle and split - idx = np.arange(len(X)); np.random.shuffle(idx) - X, y = X[idx], y[idx] - split = int(len(X) * 0.7) - - clf = LogisticRegression(solver='liblinear').fit(X[:split], y[:split]) - epsilon = 1.0 - accuracy_score(y[split:], clf.predict(X[split:]))''' accuracy_score = self._comput_adversarial_accuracy(latent1, latent2, axis=0) epsilon = 1 - accuracy_score return 2.0 * np.abs(0.5 - epsilon) - - - '''def _hook_fn(self, module, input, output): - if not self.collecting: - return - flattened_embeddings = torch.flatten(output, 1) - self._hook_features.append(flattened_embeddings.detach()) - - def register_model_hook(self, inner_model): - if hasattr(inner_model, "original_model"): - core_model = inner_model.original_model - else: - core_model = inner_model - - if hasattr(core_model, 'avgpool'): - pool_layer = core_model.avgpool - else: - pool_layer = None - for name, module in core_model.named_modules(): - if 'pool' in name: - pool_layer = module - break - - if pool_layer is None: - raise AttributeError("The target model architecture lacks an 'avgpool' layer block.") - - self.hook = pool_layer.register_forward_hook(self._hook_fn) - - def shutdown_hook(self): - if hasattr(self, 'hook') and self.hook: - self.hook.remove() - self.hook = None - ''' def _extract_attack_features(self, target_model, loader, device, target_class): - ''' - # cnahgng to black box. - target_model.eval() - all_probs = [] - all_entropies = [] - all_losses = [] - self._hook_features = [] - self.collecting = True - with torch.no_grad(): - for data, targets in loader: - data, targets = data.to(device), targets.to(device) - if target_model.__class__.__name__ == "WF_Module": - gate_signals = torch.full( - (data.size(0),), - target_class, - dtype=torch.long, - device=data.device - ) - outputs = target_model(data, target_class_indices=gate_signals) - else: - outputs = target_model(data) - - probs = torch.softmax(outputs, dim=1) - all_probs.extend(probs.cpu().numpy()) - - entropy = -torch.sum(probs * torch.log(probs + 1e-10), dim=1) - all_entropies.extend(entropy.cpu().numpy()) - - loss = self.criterion(outputs, targets) - all_losses.extend(loss.cpu().numpy()) - - self.collecting = False - - X_probs = np.array(all_probs) - X_entropy = np.array(all_entropies).reshape(-1, 1) - X_loss = np.array(all_losses).reshape(-1, 1) - X_features = np.hstack([X_probs, X_entropy, X_loss]) - if self._hook_features: - compiled_latent = torch.cat(self._hook_features, dim=0).cpu().numpy() - else: - compiled_latent = np.zeros((len(X_features), 512)) - - return X_features, compiled_latent''' target_model.eval() all_probs, all_entropies, all_losses = [], [], [] @@ -179,30 +80,7 @@ class UnlearningAttack: ]) def run_parameter_space_mia(self, unlearned_model, shadow_model, forget_loader, retain_test_loader, device, index): - '''X_shadow_mem, _ = self._extract_attack_features(shadow_model, forget_loader, device, index) - X_shadow_non, _ = self._extract_attack_features(shadow_model, retain_test_loader, device, index) - min_train = min(len(X_shadow_mem), len(X_shadow_non)) - X_train = np.vstack([X_shadow_mem[:min_train], X_shadow_non[:min_train]]) - y_train = np.concatenate([np.ones(min_train), np.zeros(min_train)]) - - attack_classifier = LogisticRegression(max_iter=1000) - attack_classifier.fit(X_train, y_train) - - X_eval_mem, latent_features = self._extract_attack_features(unlearned_model, forget_loader, device, index) - X_eval_non, retain_latent = self._extract_attack_features(unlearned_model, retain_test_loader, device, index) - - min_test = min(len(X_eval_mem), len(X_eval_non)) - X_test = np.vstack([X_eval_mem[:min_test], X_eval_non[:min_test]]) - y_test = np.concatenate([np.ones(min_test), np.zeros(min_test)]) - - predictions = attack_classifier.predict(X_test) - mia_accuracy = accuracy_score(y_test, predictions) - - #clean_centroid = np.mean(retain_latent, axis=0) - #forget_distances = np.linalg.norm(latent_features - clean_centroid, axis=1) - - return mia_accuracy, np.mean(forget_distances)''' X_shadow_mem = self._extract_attack_features(shadow_model, forget_loader, device, index) X_shadow_non = self._extract_attack_features(shadow_model, retain_test_loader, device, index) @@ -226,7 +104,7 @@ class UnlearningAttack: mia_accuracy = accuracy_score(y_test, predictions) # Note: Latent distance is removed as it's not a black-box metric - return mia_accuracy, 0.0 + return mia_accuracy def _comput_adversarial_accuracy(self, filtered, naive, axis=0): @@ -301,10 +179,8 @@ class UnlearningAttack: with open(current_log_file, "w") as f: f.write("target_class, parameter_mia_accuracy, latent_distance_tell, lookalike_accuracy, A-Dist, JS-Dist\n") - #self.register_model_hook(unlearned_instance.model) - # 1. Parameter-Space MIA and Latent Distance - parameter_mia_acc, latent_dist = self.run_parameter_space_mia( + parameter_mia_acc = self.run_parameter_space_mia( unlearned_model=unlearned_instance.model, shadow_model=base_shadow_instance.model, forget_loader=forget_loader, @@ -335,30 +211,24 @@ class UnlearningAttack: target_class=target_class ) - # 1. Calculate JS-Dist (Logit-space probability comparison) + # Calculate JS-Dist js_dist = self.calculate_js_dist(unlearned_instance.model, reference_model_torch, forget_loader, device, target_class) - # 2. Extract latent features for A-Dist - # We need features from both Unlearned and Retrained model - #_, unlearned_latent = self._extract_attack_features(unlearned_instance.model, forget_loader, device, target_class) - #_, retrained_latent = self._extract_attack_features(reference_model_torch, forget_loader, device, target_class) - # Extract features (now just one returned object) + # Extract features unlearned_features = self._extract_attack_features(unlearned_instance.model, forget_loader, device, target_class) retrained_features = self._extract_attack_features(reference_model_torch, forget_loader, device, target_class) # Calculate A-Dist using these features a_dist = self.calculate_a_dist(unlearned_features, retrained_features) - # 3. Calculate A-Dist (Replacing latent_distance) - #a_dist = self.calculate_a_dist(unlearned_latent, retrained_latent) - print(f"[{framework_name}] Class {target_class} | Parameter MIA: {parameter_mia_acc:.4f} | Latent Dist: {latent_dist:.4f} | Lookalike: {lookalike_acc:.4f}" ) + print(f"[{framework_name}] Class {target_class} | Parameter MIA: {parameter_mia_acc:.4f} Lookalike: {lookalike_acc:.4f}" ) with open(current_log_file, "a") as f: - f.write(f"{target_class},{parameter_mia_acc:.6f},{latent_dist:.6f},{lookalike_acc:.6f}, {a_dist:.6f}, {js_dist:.6f}\n") + f.write(f"{target_class},{parameter_mia_acc:.6f},{0.00000},{lookalike_acc:.6f}, {a_dist:.6f}, {js_dist:.6f}\n") return { "parameter_mia_accuracy": parameter_mia_acc, - "latent_distance": latent_dist, + "latent_distance": a_dist, "lookalike_accuracy": lookalike_acc } \ No newline at end of file diff --git a/unlearning/LinearFiltration.py b/unlearning/LinearFiltration.py index d5df62e..1b2d44a 100644 --- a/unlearning/LinearFiltration.py +++ b/unlearning/LinearFiltration.py @@ -12,10 +12,6 @@ class LinearFiltration(Strategy): def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module: model.eval() - # Freeze internal params - #for param in model.parameters(): - #param.requires_grad = False - device = next(model.parameters()).device return self.normalise(