diff --git a/Tune_new.py b/Tune_new.py index e9152f1..0145b9a 100644 --- a/Tune_new.py +++ b/Tune_new.py @@ -9,11 +9,9 @@ import Util from sets.Data import * from sets.IdentitySubset import IdentitySubset from architectures.Model import Model, Architecture -from unlearning.CertifiedRemoval import CertifiedRemoval from unlearning.CertifiedUnlearning import CertifiedUnlearning from unlearning.LinearFiltration import LinearFiltration from unlearning.WeightFiltration import WeightFiltration -from unlearning.WF import WeightF # Global Hyperparameters @@ -140,40 +138,12 @@ def run_unlearning_and_strategy_eval(env_dict, forget_class_idx, strategy, evalu train_data = env_dict["train_data"] test_data = env_dict["test_data"] - # testing valuse * * - #--------------------------------------------------------------------------- - # S1 50 5 5 5 5 5 - # S2 1000 200 1000 500 200 300 - # BS 5 5 5 5 5 5 - # scale 2000 500 8000 5000 10000 8000 - # std 0.00001 0.00001 0.00001 0.00001 0.00001 0.00001 - - # Initialize the strategy hyperparameters matching standard settings - # increase s2, decrease scale ---sweet spot - '''certified_removal = CertifiedRemoval( - target_class_index=forget_class_idx, - s1=4, - s2=350, # 350 best - unlearn_bs=5, - scale=6000.0, # 6000 was good - std=0.00001 - )''' - '''certified_removal = CertifiedUnlearning( - target_class_index=0, - l2_reg=0.0005, - gamma=0.1, - scale=7000.0, - s1=2, - s2=350, - std=1e-5, - unlearn_bs=2 - )''' # Segment specific unlearning loaders using class index boundaries - forget_train_loader, retain_train_loader = get_unlearning_loaders( + retain_train_loader , forget_train_loader= get_unlearning_loaders( dataset=train_data, forget_class_idx=forget_class_idx, batch_size=BATCH_SIZE ) - forget_test_loader, retain_test_loader = get_unlearning_loaders( + retain_test_loader, forget_test_loader = get_unlearning_loaders( dataset=test_data, forget_class_idx=forget_class_idx, batch_size=BATCH_SIZE ) @@ -189,9 +159,16 @@ def run_unlearning_and_strategy_eval(env_dict, forget_class_idx, strategy, evalu print("fine tunned model loaded into evaluation sandbox") # Execute strategic parameter unlearning step - strategy.apply(reloaded.model, forget_train_loader, retain_train_loader) + unlearned = strategy.apply(reloaded.model, train_data) strategy_in_use = strategy.__class__.__name__ + if isinstance(unlearned,nn.Module): + reloaded.model = unlearned + else: + reloaded = unlearned + + + # Define validation tracking steps dynamically evaluation_domains = [ {"loader": retain_test_loader, "mode": "retain", "label": "\n--- Performance on Retained Classes"}, @@ -215,66 +192,63 @@ def run_unlearning_and_strategy_eval(env_dict, forget_class_idx, strategy, evalu # entry if __name__ == "__main__": - # Run Data Infrastructure and Architecture Builder - runtime_environment = prepare_data_and_model_environment() - - # Baseline Evaluation - finetuning = False - # switch finetuning for tests on strategies only - run_finetuning_or_baseline_eval(runtime_environment, run_training=finetuning) - - finetuning = True - # Unlearning Iterations - for i in range(0, 1): - # strategies - # - #certified_removal = CertifiedRemoval( - # target_class_index=i, - # s1=4, - # s2=350, # 350 best - # unlearn_bs=5, - # scale=6000.0, # 6000 was good - # std=0.00009 - # ) + try: + # Run Data Infrastructure and Architecture Builder + runtime_environment = prepare_data_and_model_environment() - - + # Baseline Evaluation + finetuning = False + # switch finetuning for tests on strategies only + run_finetuning_or_baseline_eval(runtime_environment, run_training = finetuning) + + + # strategies certified_unlearning = CertifiedUnlearning( - target_class_index=i, + target_class_index=0, l2_reg=0.000002, gamma=0.1, - scale= 20000,# 16400.0, # took ages to reach this sweet spot + scale= 16400.0,# 16400.0, # took ages to reach this sweet spot s1=2, s2=300, std=0.00001, - unlearn_bs=16 + unlearn_bs=8 ) # works perfectly linear_filtration = LinearFiltration( - target_class_index=i + target_class_index=0 ) - weight_filtration = WeightF( #WeightFiltration( - target_class_index=i, - epochs=3, - lr=0.05, - gamma=5 + weight_filtration = WeightFiltration( + target_class_index=0, + epochs=6, + lr=150.0, + gamma=0.001 ) strategies = [ - # certified_unlearning, + certified_unlearning, weight_filtration, - # linear_filtration + linear_filtration ] + # Unlearning Iteration + for i in range(0, CLASS_SIZE): - print(f"\n>>> Executing Unlearning Framework for Target Identity Index: {i} <<<") - for strategy in strategies: - run_unlearning_and_strategy_eval( - runtime_environment, - forget_class_idx=i, - strategy=strategy, - evaluate= not finetuning - ) + for strategy in strategies: + + # update target class to be unlearned + strategy.set_target_class(i) + print(f"Unlearning class {i} with {strategy.strategy_name}") + + # forget + run_unlearning_and_strategy_eval( + runtime_environment, + forget_class_idx=i, + strategy=strategy, + evaluate = not finetuning + ) + + except KeyboardInterrupt: + print("program interrupted. Exit!") diff --git a/Util.py b/Util.py index 66d302f..ee84cb0 100644 --- a/Util.py +++ b/Util.py @@ -45,4 +45,7 @@ def _initialize_log_file(log_file): def log_metric(log_file, execution_time: float): """Appends the execution time to this strategy's specific file.""" with open(log_file, "a") as f: - f.write(f"{execution_time:.6f}\n") \ No newline at end of file + f.write(f"{execution_time:.6f}\n") + + + diff --git a/architectures/Model.py b/architectures/Model.py index 287be1e..853f221 100644 --- a/architectures/Model.py +++ b/architectures/Model.py @@ -7,7 +7,7 @@ import time import numpy as np from sklearn.metrics import classification_report from pathlib import Path -from unlearning.Strategy import Strategy +#from unlearning.Strategy import Strategy import copy from torch.optim.lr_scheduler import CosineAnnealingLR @@ -84,7 +84,7 @@ class Model(ABC): print(f'Model loaded from {file_path}') - def unlearn(self, strategy: Strategy, forget_loader, retain_loader): + def unlearn(self, strategy: 'Strategy', forget_loader, retain_loader): """ Executes a targeted unlearning strategy and profiles efficiency """ print(f"Executing: {strategy.__class__.__name__}...") @@ -103,6 +103,7 @@ class Model(ABC): Evaluates the model, prints terminal reports, and routes metrics to a file logger based on the current context mode. """ + self.model.eval() all_preds, all_labels = [], [] print(f"\nEvaluating Domain: [{mode}]...") diff --git a/sets/Data.py b/sets/Data.py index 36138a7..ca2ff7e 100644 --- a/sets/Data.py +++ b/sets/Data.py @@ -1,5 +1,5 @@ from torchvision import datasets, transforms -from torch.utils.data import Dataset, DataLoader, Subset +from torch.utils.data import Dataset, DataLoader, Subset, ConcatDataset import torch import numpy as np import os @@ -181,4 +181,59 @@ def get_unlearning_loaders(dataset: Dataset, forget_class_idx: int, batch_size: print(f"[Data Split] Local Class {forget_class_idx}: {len(forget_subset)} samples | Remaining Classes: {len(retain_subset)} samples.") - return forget_loader, retain_loader \ No newline at end of file + return retain_loader, forget_loader + + +def vertical_split(dataset, batch_size,num_classes): + """ + Executes a class-wise vertical split. + Divides the samples of every single identity class exactly in half: + 50% of each class goes to the Retain Set, 50% goes to the Forget Set. + """ + + # 1. Group dataset indices by their respective ground-truth classes + class_to_indices = {c: [] for c in range(num_classes)} + + print(" [Vertical Split] Tracking class indices across the combined dataset...") + for idx in range(len(dataset)): + # Extract the label cleanly from the underlying dataset structure + _, label = dataset[idx] + if label in class_to_indices: + class_to_indices[label].append(idx) + + retain_indices = [] + forget_indices = [] + + # 2. Slice each class identity vertically (exactly 50/50) + for c, indices in class_to_indices.items(): + if len(indices) < 2: + print(f" Warning: Class {c} has fewer than 2 samples. Cannot split vertically.") + retain_indices.extend(indices) + continue + + # Deterministic shuffle per class to ensure honest distribution before splitting + np.random.shuffle(indices) + + mid = len(indices) // 2 + forget_indices.extend(indices[:mid]) # First half assigned to unlearning + retain_indices.extend(indices[mid:]) # Second half assigned to retention + + print(f" Vertical split complete: Retain Index Size = {len(retain_indices)} | Forget Index Size = {len(forget_indices)}") + + # 3. Construct lightweight PyTorch Subsets using our sliced index maps + retain_subset = Subset(dataset, retain_indices) + forget_subset = Subset(dataset, forget_indices) + + # 4. Return pristine, shuffled DataLoaders mirroring your environment's batch specifications + retain_loader = DataLoader(retain_subset, batch_size=batch_size, shuffle=True) + forget_loader = DataLoader(forget_subset, batch_size=batch_size, shuffle=True) + + return retain_loader, forget_loader + +def _combine_set(loader_one, loader_two): + full_train_dataset = ConcatDataset([loader_one.dataset, loader_two.dataset]) + return DataLoader( + full_train_dataset, + batch_size=loader_one.batch_size, + shuffle=True + ) \ No newline at end of file diff --git a/unlearning/CertifiedRemoval.py b/unlearning/CertifiedRemoval.py deleted file mode 100644 index 058fd22..0000000 --- a/unlearning/CertifiedRemoval.py +++ /dev/null @@ -1,214 +0,0 @@ -import torch -import torch.nn as nn -from torch.utils.data import DataLoader, RandomSampler -from torch.autograd import grad -from unlearning.Strategy import Strategy - -class CertifiedRemoval(Strategy): - """ - Implements Certified Unlearning for non-convex DNNs (Zhang et al.). - Uses a modified, stabilized stochastic Newton step using Taylor-expansion - HVP estimation across the entire parameter space, capped with calibrated noise. - """ - def __init__(self, target_class_index: int, l2_reg: float = 0.0005, - gamma: float = 0.01, scale: float = 1000.0, - s1: int = 10, s2: int = 1000, std: float = 0.001, unlearn_bs: int = 2): - super().__init__(target_class_index) - self.l2_reg = l2_reg - self.gamma = gamma - self.scale = scale - self.s1 = s1 - self.s2 = s2 - self.std = std - self.unlearn_bs = unlearn_bs - - ''' - def _compute_loss_gradient(self, model, loader, device: torch.device): - model.eval() - criterion = nn.CrossEntropyLoss(reduction='sum') - params = [p for p in model.parameters() if p.requires_grad] - grad_accumulator = [torch.zeros_like(p).cpu() for p in params] - total_samples = 0 - - for data, targets in loader: - total_samples += targets.shape[0] - data, targets = data.to(device), targets.to(device) - outputs = model(data) - - mini_grads = list(grad(criterion(outputs, targets), params)) - for i in range(len(grad_accumulator)): - grad_accumulator[i] += mini_grads[i].cpu().detach() - - for i in range(len(grad_accumulator)): - grad_accumulator[i] /= total_samples - - l2_reg_term = 0.0 - for param in model.parameters(): - l2_reg_term += torch.norm(param, p=2) - - reg_grads = list(grad(self.l2_reg * l2_reg_term, params)) - for i in range(len(grad_accumulator)): - grad_accumulator[i] += reg_grads[i].cpu().detach() - - return [p.to(device) for p in grad_accumulator]''' - def _compute_loss_gradient(self, model, loader, device: torch.device): - model.eval() - # Use reduction='sum' matching the original framework - criterion = nn.CrossEntropyLoss(reduction='sum') - params = [p for p in model.parameters() if p.requires_grad] - grad_accumulator = [torch.zeros_like(p).cpu() for p in params] - total_samples = 0 - - for data, targets in loader: - total_samples += targets.shape[0] - data, targets = data.to(device), targets.to(device) - outputs = model(data) - - loss = criterion(outputs, targets) - - # Incorporate L2 weight regularization directly inside the backprop graph - # to keep scaling bounded and aligned with the data volume - l2_reg_term = 0.0 - for param in model.parameters(): - if param.requires_grad: - l2_reg_term += torch.norm(param, p=2) - - total_loss = loss + (self.l2_reg * l2_reg_term) - - mini_grads = list(grad(total_loss, params, retain_graph=False)) - for i in range(len(grad_accumulator)): - grad_accumulator[i] += mini_grads[i].cpu().detach() - - for i in range(len(grad_accumulator)): - grad_accumulator[i] /= total_samples - - return [p.to(device) for p in grad_accumulator] - - - def grad_batch(batch_loader, lam, model, device): - model.eval() - criterion = nn.CrossEntropyLoss(reduction='sum') - params = [p for p in model.parameters() if p.requires_grad] - grad_batch = [torch.zeros_like(p).cpu() for p in params] - num = 0 - for batch_idx, (data, targets) in enumerate(batch_loader): - num += targets.shape[0] - data, targets = data.to(device), targets.to(device) - outputs = model(data) - - grad_mini = list(grad(criterion(outputs, targets), params)) - for i in range(len(grad_batch)): - grad_batch[i] += grad_mini[i].cpu().detach() - - for i in range(len(grad_batch)): - grad_batch[i] /= num - - l2_reg = 0 - for param in model.parameters(): - l2_reg += torch.norm(param, p=2) - grad_reg = list(grad(lam * l2_reg, params)) - for i in range(len(grad_batch)): - grad_batch[i] += grad_reg[i].cpu().detach() - return [p.to(device) for p in grad_batch] - - def _hvp(self, loss, params, v): - first_grads = grad(loss, params, retain_graph=True, create_graph=True) - elemwise_products = 0 - for grad_elem, v_elem in zip(first_grads, v): - elemwise_products += torch.sum(grad_elem * v_elem) - # FIX 1: Set create_graph to False to prevent massive nested graph accumulation - return grad(elemwise_products, params, create_graph=False) - - def _stochastic_newton_update(self, g, retain_dataset, model, device): - model.eval() - criterion = nn.CrossEntropyLoss() - params = [p for p in model.parameters() if p.requires_grad] - h_res = [torch.zeros_like(p) for p in g] - - for _ in range(self.s1): - h_estimate = [p.clone() for p in g] - sampler = RandomSampler(retain_dataset, replacement=True, num_samples=self.unlearn_bs * self.s2) - res_loader = DataLoader(retain_dataset, batch_size=self.unlearn_bs, sampler=sampler) - res_iter = iter(res_loader) - - for j in range(self.s2): - try: - data, target = next(res_iter) - except StopIteration: - res_iter = iter(res_loader) - data, target = next(res_iter) - - data, target = data.to(device), target.to(device) - outputs = model(data) - - loss = criterion(outputs, target) - l2_reg_term = 0.0 - for param in model.parameters(): - l2_reg_term += torch.norm(param, p=2) - loss += (self.l2_reg + self.gamma) * l2_reg_term - - h_s = self._hvp(loss, params, h_estimate) - - with torch.no_grad(): - for k in range(len(params)): - # FIX 2: Added .detach() to decouple history strings across iterative update blocks - #h_estimate[k] = (h_estimate[k] + g[k] - h_s[k] / self.scale).detach() - next_estimate = h_estimate[k].data + g[k].data - (h_s[k].data / self.scale) - h_estimate[k] = next_estimate.clone() - del h_s, loss, outputs - - for k in range(len(params)): - h_res[k] = h_res[k] + h_estimate[k] / self.scale - - return [p / self.s1 for p in h_res] - - '''def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module: - device = next(model.parameters()).device - - num_forget = len(forget_loader.dataset) - num_retain = len(retain_loader.dataset) - scaling_ratio = num_forget / num_retain - - print(">> Calculating base gradients over target FORGET set...") - # FIX 3: Base gradients MUST be evaluated from forget_loader to drop target class distributions - g = self._compute_loss_gradient(model, forget_loader, device) - - print(">> Estimating non-convex inverse Hessian trajectories via Taylor series...") - retain_dataset = retain_loader.dataset - delta = self._stochastic_newton_update(g, retain_dataset, model, device) - - print(">> Applying stabilized parameter adjustments and randomized certification noise...") - with torch.no_grad(): - for i, param in enumerate(model.parameters()): - if param.requires_grad: - noise = self.std * torch.randn(param.data.size(), device=device) - #param.data.add_(-delta[i] + noise) - param.data.add_(scaling_ratio * delta[i] + noise) - - print(">> Certified Unlearning process completed successfully across the complete landscape.") - return model''' - def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module: - device = next(model.parameters()).device - - print(">> Calculating stable base gradients over the RETAIN set...") - # To match the author's snippet perfectly, g MUST be computed on the retain data. - # If this loader is too large for your VRAM, use a smaller batch size (e.g. 16 or 32) - # in your main training script when creating retain_loader. - g = self._compute_loss_gradient(model, retain_loader, device) - - print(">> Estimating non-convex inverse Hessian trajectories via Taylor series...") - retain_dataset = retain_loader.dataset - delta = self._stochastic_newton_update(g, retain_dataset, model, device) - - print(">> Applying parameter removal adjustments (-delta)...") - with torch.no_grad(): - for i, param in enumerate(model.parameters()): - if param.requires_grad: - noise = self.std * torch.randn(param.data.size(), device=device) - - # MATCHING THE SNIPPET: Subtract delta exactly as the authors do - # This removes the influence trace of the omitted data. - param.data.add_(-delta[i] + noise) - - print(">> Certified Unlearning process completed successfully.") - return model \ No newline at end of file diff --git a/unlearning/Certified_facebook.py b/unlearning/Certified_facebook.py deleted file mode 100644 index 0b8bcd1..0000000 --- a/unlearning/Certified_facebook.py +++ /dev/null @@ -1,123 +0,0 @@ -import torch -import torch.nn as nn -import math -from torch.utils.data import DataLoader -from unlearning.Strategy import Strategy - -class CertifiedRemovalFacebook(Strategy): - """ - Implements Certified Removal (Guo et al.) mapped for Multi-Class models - by executing a single-class One-vs-Rest (OvR) block-removal update step. - Math matches the facebookresearch/certified-removal reference repository. - """ - def __init__(self, target_class_index: int, removal_bound: float, epsilon: float, l2_reg: float = 0.1): - super().__init__(target_class_index=target_class_index) - self.removal_bound = removal_bound # gamma in the paper - self.epsilon = epsilon # Privacy budget - self.l2_reg = l2_reg # Lambda (regularization term) - - def _get_features(self, backbone: nn.Module, loader: DataLoader, device: torch.device): - """Passes data through the frozen ResNet backbone to extract embedding features.""" - backbone.eval() - all_features = [] - - with torch.no_grad(): - for inputs, _ in loader: - inputs = inputs.to(device) - # Pass through frozen backbone to get the 2048-dimensional embedding - features = backbone(inputs) - all_features.append(features.cpu()) - - return torch.cat(all_features, dim=0) - - def _fb_lr_grad(self, w, X, y, lam): - """ - Replicates exact lr_grad calculation from Facebook's codebase. - Note: The resulting gradient has a flipped sign due to the structure of (z - 1). - """ - # X.mv(w) computes raw linear margins - z = torch.sigmoid(y * X.mv(w)) - # Gradient formula: X^T * ((z - 1) * y) + lambda * N * w - return X.t().mv((z - 1) * y) + lam * X.size(0) * w - - def _fb_lr_hessian_inv(self, w, X, y, lam, device, batch_size=50000): - """ - Replicates exact lr_hessian_inv calculation from Facebook's codebase. - Scales the L2 regularization matrix explicitly by dataset row count (N * lambda * I). - """ - z = torch.sigmoid(X.mv(w).mul_(y)) - D = z * (1 - z) # Element-wise variance vector - - H = None - num_batch = int(math.ceil(X.size(0) / batch_size)) - for i in range(num_batch): - lower = i * batch_size - upper = min((i + 1) * batch_size, X.size(0)) - X_i = X[lower:upper] - - # Stepwise feature weighting via element-wise variance columns - if H is None: - H = X_i.t().mm(D[lower:upper].unsqueeze(1) * X_i) - else: - H += X_i.t().mm(D[lower:upper].unsqueeze(1) * X_i) - - # Scale identity buffer by dataset split size: lambda * N_retain - reg_matrix = lam * X.size(0) * torch.eye(X.size(1), device=device).float() - return torch.linalg.inv(H + reg_matrix) - - def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module: - """ - Applies Certified Removal strictly to the target class parameters - belonging to the final fully connected layer (model.fc). - """ - device = next(model.parameters()).device - k = self.target_class_index - - # Isolate final layer and extract raw deep embeddings using frozen backbone - linear_head = model.fc - model.fc = nn.Identity() - - print(">> Extracting deep features from model backbone...") - X_retain = self._get_features(model, retain_loader, device).to(device) - X_forget = self._get_features(model, forget_loader, device).to(device) - - # Restore the classification head back - model.fc = linear_head - - # Extract current model weight row for the target class channel - w_k = model.fc.weight.data[k].clone().to(device) - - # Create One-vs-Rest binary target indicator arrays (+1.0 / -1.0) - # Retain dataset instances are negative labels (-1.0) for the target class channel - y_retain_binary = torch.full((X_retain.size(0),), -1.0, device=device) - # Forget dataset instances are positive labels (+1.0) for the target class channel - y_forget_binary = torch.full((X_forget.size(0),), 1.0, device=device) - - # Compute Inverse Hessian (on Retain Data) and Gradient (on Forget Data) - H_inv = self._fb_lr_hessian_inv(w_k, X_retain, y_retain_binary, self.l2_reg, device) - grad_forget = self._fb_lr_grad(w_k, X_forget, y_forget_binary, self.l2_reg) - - # 5. Compute the Weight Update Step Vector (Delta) - multiplier = 0.5 - delta_w_k = torch.mv(H_inv, grad_forget) * multiplier - - # Verify Theoretical Removal Bound Criteria - norm_delta = torch.norm(delta_w_k).item() - if norm_delta > self.removal_bound: - print(f"!! Warning: Removal budget exceeded! Norm: {norm_delta:.4f} > Bound: {self.removal_bound}") - else: - print(f">> Certificate valid. Norm: {norm_delta:.4f} <= Bound: {self.removal_bound}") - - # Apply Update (Using '+' since Facebook's grad calculation yields a negative sign output) - new_w_k = w_k + delta_w_k - - # Calibrate and Inject Perturbation Noise (Objective Perturbation Verification) - sigma = 2.0 / (self.l2_reg * self.epsilon) - noise = torch.randn_like(new_w_k, device=device) * (sigma / X_retain.size(0)) - new_w_k = new_w_k + noise - - # Commit updated weight vector row back into model head parameters in-place - model.fc.weight.data[k] = new_w_k - - print(">> Certified Removal process completed successfully.") - return model \ No newline at end of file diff --git a/unlearning/LastK_Certified.py b/unlearning/LastK_Certified.py deleted file mode 100644 index 6aae56a..0000000 --- a/unlearning/LastK_Certified.py +++ /dev/null @@ -1,125 +0,0 @@ -import torch -import torch.nn as nn -from torch.utils.data import DataLoader -from unlearning.Strategy import Strategy - -class LastKCertifiedRemoval(Strategy): - """ - Implements Certified Removal (Guo et al.) scaled up to the last K layers - of a ResNet50 network by flattening sub-graph parameters into a convex sub-problem. - """ - def __init__(self, removal_bound: float, epsilon: float, l2_reg: float = 0.1): - super().__init__() - self.removal_bound = removal_bound - self.epsilon = epsilon - self.l2_reg = l2_reg - - def _split_model(self, model: nn.Module): - """ - Splits ResNet50 into a frozen feature backbone and an active unlearning head. - Here, 'Last K Layers' includes layer4 and the fc classification head. - """ - # Feature Backbone: Everything up to layer3 - backbone = nn.Sequential( - model.conv1, - model.bn1, - model.relu, - model.maxpool, - model.layer1, - model.layer2, - model.layer3 - ) - - # Active Head: Layer4, global pooling, and the final linear layer - unlearning_head = nn.Sequential( - model.layer4, - model.avgpool, - nn.Flatten(1), - model.fc - ) - - return backbone, unlearning_head - - def _get_intermediate_features(self, backbone: nn.Module, loader: DataLoader, device: torch.device): - """Extracts features from the exit point of the frozen backbone (post-layer3).""" - backbone.eval() - all_features = [] - all_labels = [] - - with torch.no_grad(): - for inputs, labels in loader: - inputs = inputs.to(device) - features = backbone(inputs) - all_features.append(features.cpu()) - all_labels.append(labels.cpu()) - - return torch.cat(all_features, dim=0), torch.cat(all_labels, dim=0) - - def apply(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module: - """ - Extracts intermediate features and updates the parameters of the last blocks - using the exact inverse-Hessian influence step. - """ - device = next(model.parameters()).device - - # 1. Slice the ResNet graph structural components - backbone, unlearning_head = self._split_model(model) - - print(">> Extracting intermediate structural features from layer3 exit...") - retain_feats, retain_labels = self._get_intermediate_features(backbone, retain_loader, device) - forget_feats, forget_labels = self._get_intermediate_features(backbone, forget_loader, device) - - # 2. Flatten target weights from the active head into a 1D optimization tensor - # For simplicity and mathematical stability, we isolate the final layer's weights - # inside the active head for the exact Hessian tracking step - target_layer = unlearning_head[-1] # This points straight to model.fc - w = target_layer.weight.data.clone().cpu() - - # 3. Compute Exact Hessian over intermediate embeddings - # ResNet50's layer4 expands channels to 2048, creating a 2048x2048 matrix context - print(">> Computing exact sub-graph Hessian matrix...") - N_retain = retain_feats.size(0) - - # Pool the feature maps if they haven't been flattened yet by the head module - if len(retain_feats.shape) > 2: - retain_flat = torch.mean(retain_feats, dim=[2, 3]) - forget_flat = torch.mean(forget_feats, dim=[2, 3]) - else: - retain_flat = retain_feats - forget_flat = forget_feats - - X_T_X = torch.matmul(retain_flat.t(), retain_flat) - reg_matrix = self.l2_reg * torch.eye(retain_flat.size(1)) - Hessian = (X_T_X / N_retain) + reg_matrix - - # 4. Calculate gradients relative to the forgotten target features - print(">> Calculating forget set gradients...") - num_classes = w.size(0) - forget_labels_one_hot = torch.nn.functional.one_hot(forget_labels, num_classes=num_classes).float() - - preds_forget = torch.matmul(forget_flat, w.t()) - error = preds_forget - forget_labels_one_hot - grad_forget = torch.matmul(error.t(), forget_flat) / forget_flat.size(0) - - # 5. Apply Newton Step optimization update - print(">> Inverting optimization subspace via system solver...") - try: - delta_w_t = torch.linalg.solve(Hessian, grad_forget.t()) - delta_w = delta_w_t.t() - except RuntimeError: - print(">> Warning: Subspace Hessian is singular. Using pseudo-inverse fallback.") - delta_w = torch.matmul(grad_forget, torch.linalg.pinv(Hessian).t()) - - # 6. Apply Weight Adjustment Bounds Check - new_w = w + delta_w - norm_delta = torch.norm(delta_w).item() - if norm_delta > self.removal_bound: - print(f"!! Warning: Removal budget exceeded! Norm: {norm_delta:.4f} > Bound: {self.removal_bound}") - else: - print(f">> Certificate valid. Subspace Norm: {norm_delta:.4f} <= Bound: {self.removal_bound}") - - # 7. Write weights directly back into the live ResNet50 instance - model.fc.weight.data = new_w.to(device) - - print(">> Last K Layers Certified Removal complete.") - return model \ No newline at end of file diff --git a/unlearning/LinearFiltration.py b/unlearning/LinearFiltration.py index ebf9343..1cc174c 100644 --- a/unlearning/LinearFiltration.py +++ b/unlearning/LinearFiltration.py @@ -2,6 +2,7 @@ import torch import torch.nn as nn from .Strategy import Strategy from torch.utils.data import DataLoader +from sets.Data import get_unlearning_loaders, _combine_set class LinearFiltration(Strategy): def __init__(self, target_class_index): @@ -23,40 +24,8 @@ class LinearFiltration(Strategy): forget_index=self.target_class_index ) - # FIX: Added staticmethod decorator - @staticmethod - def get_features(model, inputs): - # For ResNet, pass through everything up to the fc layer - x = model.conv1(inputs) - x = model.bn1(x) - x = model.relu(x) - x = model.maxpool(x) - - x = model.layer1(x) - x = model.layer2(x) - x = model.layer3(x) - x = model.layer4(x) - - x = model.avgpool(x) - x = torch.flatten(x, 1) - return x - - @staticmethod - def _calculate_filtration_matrix(num_classes: int, forget_class: int, device: torch.device) -> torch.Tensor: - A = torch.eye(num_classes, device=device) - num_remaining = num_classes - 1 - - for j in range(num_classes): - if j == forget_class: - A[forget_class, j] = 0.0 - else: - A[forget_class, j] = 1.0 / num_remaining - - return A - - @staticmethod - def _sums_and_counts(model, num_classes, retain_loader, forget_loader, device, forget_index, h_dim): + def _sums_and_counts(self, model, num_classes, loader, device, forget_index, h_dim): model.eval() sums = torch.zeros(num_classes, h_dim, device=device) @@ -64,11 +33,11 @@ class LinearFiltration(Strategy): # Generate values for retain with torch.no_grad(): - for inputs, targets in retain_loader: + for inputs, targets in loader: inputs = inputs.to(device) targets = targets.to(device) - # FIX: Call get_features instead of model() directly - outputs = LinearFiltration.get_features(model, inputs) + # predictions + outputs = model(inputs) for j in range(num_classes): if j == forget_index: @@ -79,65 +48,54 @@ class LinearFiltration(Strategy): sums[j] += outputs[mask].sum(dim=0) counts[j] += mask.sum() - # Values for forget - with torch.no_grad(): - for inputs, targets in forget_loader: - inputs = inputs.to(device) - targets = targets.to(device) - # FIX: Call get_features instead of model() directly - outputs = LinearFiltration.get_features(model, inputs) - - mask = (targets == forget_index) - - if mask.any(): - sums[forget_index] += outputs[mask].sum(dim=0) - counts[forget_index] += mask.sum() - return sums, counts - @staticmethod - def _get_means(model, num_classes, retain_loader, forget_loader, device, forget_index): - h_dim = model.fc.in_features - sums, counts = LinearFiltration._sums_and_counts( + # + def _get_means(self,model, num_classes, loader, device, forget_index): + h_dim = model.fc.out_features + + # all predictions + sums, counts = self._sums_and_counts( model=model, num_classes=num_classes, - retain_loader=retain_loader, - forget_loader=forget_loader, + loader=loader, device=device, forget_index=forget_index, h_dim=h_dim ) - A = [] - for i in range(num_classes): - if counts[i] > 0: - A.append(sums[i] / counts[i]) - else: - A.append(torch.zeros(h_dim, device=device)) - - # CORRECT: Stack along dim=0 to make it (num_classes, h_dim) - return torch.stack(A, dim=0) + #A = [] + + counts_safe = counts.unsqueeze(1) + A = torch.where( + counts_safe > 0, + sums / counts_safe, + torch.zeros_like(sums) + ) + # 6 + return A - @staticmethod - def _compute_z(tensor, forget_index): - # Now tensor has shape (num_classes, h_dim) -> tensor.shape[0] is num_classes + # 9 + def _compute_z(self, tensor, forget_index): + K = tensor.shape[0] - # pi_a0 should match the feature space dimensions (h_dim) - pi_a0 = torch.zeros(tensor.shape[1], device=tensor.device) + # pi_a_forget should match the feature space dimensions (h_dim) + pi_a_f = torch.zeros(tensor.shape[1], device=tensor.device) - t_1 = pi_a0 - a0 = tensor[forget_index, :] # Extracting the row vector for the forgotten class + t_1 = pi_a_f + # Extracting the row vector for the forgotten class + a_f = tensor[forget_index, :] - mask_a0 = torch.ones( - a0.shape[0], + mask_a_f = torch.ones( + a_f.shape[0], dtype=torch.bool, device=tensor.device ) # We compute the target shift over features - t_2 = -(1.0 / (K - 1)) * a0[mask_a0].sum() + t_2 = -(1.0 / (K - 1)) * a_f[mask_a_f].sum() mask_rows = torch.ones(K, dtype=torch.bool, device=tensor.device) mask_rows[forget_index] = False @@ -148,21 +106,23 @@ class LinearFiltration(Strategy): return t_1 + t_2 + t_3 - @staticmethod - def normalise(model, retain_loader, forget_loader, device, forget_index): + # Normalisation filtration + def normalise(self, model, retain_loader, forget_loader, device, forget_index): W = model.fc.weight.data.clone() num_classes = W.shape[0] - - A = LinearFiltration._get_means( + + # we combine the data so we can calculate the mean of prdictions + full_loader = _combine_set(retain_loader, forget_loader) + # 8 + A = self._get_means( model=model, num_classes=num_classes, - retain_loader=retain_loader, - forget_loader=forget_loader, + loader=full_loader, device=device, forget_index=forget_index ) - - Z = LinearFiltration._compute_z(tensor=A, forget_index=forget_index) + # 9 + Z = self._compute_z(tensor=A, forget_index=forget_index) B_Z_rows = [] for i in range(num_classes): @@ -172,13 +132,24 @@ class LinearFiltration(Strategy): # Retained classes maintain their original ideal feature directions B_Z_rows.append(A[i]) + # 10 # Stack back along dim=0 to match (num_classes, h_dim) + # to get mean B_Z = torch.stack(B_Z_rows, dim=0) A_inv = torch.linalg.pinv(A) - + # 11 W_Z = B_Z @ A_inv @ W + # 12 model.fc.weight.copy_(W_Z) - return model \ No newline at end of file + return model + + # overriden function + def _split_data(self, dataset): + return get_unlearning_loaders( + dataset=dataset, + forget_class_idx=self.target_class_index, + batch_size = 32 + ) \ No newline at end of file diff --git a/unlearning/Strategy.py b/unlearning/Strategy.py index 1e29255..4e03465 100644 --- a/unlearning/Strategy.py +++ b/unlearning/Strategy.py @@ -16,12 +16,20 @@ class Strategy: self.log_file = Path(f"reports/{self.strategy_name}/metrics.txt") Util._initialize_log_file(log_file= self.log_file) - def apply(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module: + def set_target_class(self, target_class_index: int): + """Dynamically switch the unlearning target without retraining.""" + self.target_class_index = target_class_index + + + def apply(self, model: nn.Module, dataset) -> nn.Module: """ Wraps the unlearning execution with automated timing and strategy-specific logging. DO NOT override this method in subclasses. Override _run instead. """ start_time = time.perf_counter() + + + retain_loader, forget_loader = self._split_data(dataset) # Execute core unlearning logic processed_model = self._run(model, forget_loader, retain_loader) @@ -41,4 +49,12 @@ class Strategy: def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module: """Subclasses implement their core unlearning logic here.""" - raise NotImplementedError \ No newline at end of file + raise NotImplementedError + + ''' + different strategies split data in to different partitions differently. + So a strategy will implement its own and since this part is startegy specific. + not all should compute it the same. + ''' + def _split_data(self,dataset): + pass \ No newline at end of file diff --git a/unlearning/WeightFiltration.py b/unlearning/WeightFiltration.py index 48a321e..9678b5c 100644 --- a/unlearning/WeightFiltration.py +++ b/unlearning/WeightFiltration.py @@ -1,126 +1,135 @@ import torch import torch.nn as nn import torch.optim as optim -from torch.utils.data import DataLoader +from torch.utils.data import DataLoader, ConcatDataset, Subset from unlearning.Strategy import Strategy -from .wf.WF_Net import WF_Net +import numpy as np +from sklearn.metrics import classification_report +from architectures.WFNet import WF_Net_Model + +from sets.Data import vertical_split class WeightFiltration(Strategy): - """ - Verbatim implementation of Poppi et al.'s WF-Net framework. - Directly filters the convolutional weights of a target layer using a learnable - channel mask, optimizing it via weight-space regularization. - """ - def __init__(self, target_class_index: int, epochs: int = 10, lr: float = 0.2, gamma: float = 10.0): + def __init__(self, + target_class_index: int, + num_classes: int = 20, + epochs: int = 6, + lr: float = 100.0, + gamma: float = 0.01, + ): super().__init__(target_class_index=target_class_index) self.epochs = epochs self.lr = lr self.gamma = gamma - #self.alpha = None + self.num_classes = num_classes + self.wf_model = None + self.lambda_1 = 25 + + + def _optimise_filter(self, model: nn.Module, retain_loader: DataLoader, forget_loader: DataLoader, device) -> nn.Module: + # new WF_Model instance + wf_model = WF_Net_Model( + device=device, + size=self.num_classes, + original_model=model, + target_class_index=self.target_class_index + ) - - - - def _optimise_filter(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader, device): - # 1. Initialize the wrapper with your pre-trained model - num_classes = model.fc.out_features - wf_model = WF_Net(original_model=model, num_classes=num_classes).to(device) - - # 2. ONLY optimize alpha (everything else is frozen inside the wrapper) - optimizer = optim.Adam([wf_model.alpha], lr=self.lr) + # a WF_net module to be trained (unlearned) to generate alpha + wf_net = wf_model.get() + optimizer = optim.SGD([wf_net.alpha], lr=self.lr) criterion = nn.CrossEntropyLoss() for epoch in range(self.epochs): - forget_iter = iter(forget_loader) t_loss_r, t_loss_f = 0.0, 0.0 steps = 0 - for r_inputs, r_labels in retain_loader: + # forget and retain + for (r_inputs, r_labels), (f_inputs, f_labels) in zip(retain_loader, forget_loader): r_inputs, r_labels = r_inputs.to(device), r_labels.to(device) - - # Pull the matching forget batch input - try: - f_inputs, _ = next(forget_iter) - except StopIteration: - forget_iter = iter(forget_loader) - f_inputs, _ = next(forget_iter) - f_inputs = f_inputs.to(device) + f_inputs, f_labels = f_inputs.to(device), f_labels.to(device) optimizer.zero_grad() - # --- APPLY ALGORITHM 1 FORWARD PASS TO BOTH INPUTS --- - # Pass the input batch AND the target unlearn class index - outputs_r = wf_model(r_inputs, target_unlearn_class=self.target_class_index) - outputs_f = wf_model(f_inputs, target_unlearn_class=self.target_class_index) + # retain data paired with randomly selected rows of alpha to compute the retaining loss + random_rows = [] + for label in r_labels: + allowed = [i for i in range(self.num_classes) if i != label.item()] + random_rows.append(np.random.choice(allowed)) - # Compute Losses using Poppi et al.'s temperature scaled entropy + gate_signals_r = torch.tensor(random_rows, dtype=torch.long, device=device) + outputs_r = wf_net(r_inputs, target_class_indices=gate_signals_r) loss_r = criterion(outputs_r, r_labels) + + # Forget set is paired with corresponding labels as row selectors for alpha + # and used to compute unlearning loss + outputs_f = wf_net(f_inputs, target_class_indices=f_labels) - temperature = 3.0 - logits_f_scaled = outputs_f / temperature + loss_f = 0.0 + classes_in_batch = 0 - # Compute uniform target entropy per-sample, then average over the batch - log_probs_f = torch.log_softmax(logits_f_scaled, dim=-1) - uniform_target = torch.ones_like(logits_f_scaled) / num_classes - loss_f = -torch.sum(uniform_target * log_probs_f, dim=-1).mean() + # every image of class c will unlearn over the same row of alpha_l (poppi et al page 5) + for c in range(self.num_classes): + class_mask = (f_labels == c) + if not class_mask.any(): + continue + + labels_c = f_labels[class_mask] + + # Slice the existing outputs instead of recalculating a forward pass + outputs_f_c = outputs_f[class_mask] + + loss_f_ce = criterion(outputs_f_c, labels_c) + + # Poppi et al. suggest employing reciprocal of the forget loss + # to avoid shortcomings of negative gradient approach + loss_f += 1.0 / (loss_f_ce + 1e-6) + classes_in_batch += 1 + + # Average forget loss by number of distinct classes seen in this batch + if classes_in_batch > 0: + loss_f = loss_f / classes_in_batch + + # Regilarisation penality + loss_reg = torch.sum(1.0 - torch.sigmoid(wf_net.alpha)) - total_loss = loss_r + (self.gamma * loss_f) + # back propagation + total_loss = loss_r + (self.lambda_1 * loss_f) + (self.gamma * loss_reg) total_loss.backward() optimizer.step() t_loss_r += loss_r.item() - t_loss_f += loss_f.item() + t_loss_f += loss_f.item() if classes_in_batch > 0 else 0.0 steps += 1 print(f" Epoch {epoch+1}/{self.epochs} | Retain Loss: {t_loss_r/steps:.4f} | Forget Loss: {t_loss_f/steps:.4f}") + return wf_model - - def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module: device = next(model.parameters()).device - model.eval() + model.eval() - # In WF-Net, the mask targets the last major convolutional block - # For ResNet-18, that is the final conv layer in layer4 block 1 - if hasattr(model, 'layer4') and len(model.layer4) > 1: - target_conv = model.layer4[1].conv2 + if self.wf_model is None: + print(">> Initializing and compiling global WF-Net matrix (Run Once for all classes)...") + + self.wf_model = self._optimise_filter( + model, + retain_loader=retain_loader, + forget_loader=forget_loader, + device=device + ) else: - raise AttributeError("Model architecture does not match expected ResNet-18 structure.") + print(f">> Gating matrix loaded. Switching layout to target class index: {self.target_class_index}") + self.wf_model.target_class_index = self.target_class_index - # Store a pristine, non-grad copy of the original trained weights - # Shape of conv2.weight: (out_channels, in_channels, kernel_size, kernel_size) -> e.g., (512, 512, 3, 3) - original_weights = target_conv.weight.data.clone().detach() - out_channels = original_weights.shape[0] - - # Initialize alpha gate vector matching Poppi et al.'s initialization range - # Shape: (out_channels,) -> acting directly as a filter-level gate - #self.alpha = nn.Parameter(torch.ones(out_channels, device=device) * 1.5) - - # Freeze the global model graph; only optimize our filter parameter mask - for p in model.parameters(): - p.requires_grad = False - #self.alpha.requires_grad = True - - wf_model = self._optimise_filter( - model, - forget_loader=forget_loader, - retain_loader=retain_loader, - device=device, + return self.wf_model + + def _split_data(self, dataset): + return vertical_split( + dataset= dataset, + batch_size=32, + num_classes=self.num_classes ) - - # --- PERMANENT BAKING STEP --- - with torch.no_grad(): - # Grab the alpha mask vector for the forgotten class and cast to 4D tensor shape - final_mask = torch.sigmoid(wf_model.alpha[self.target_class_index]).view(-1, 1, 1, 1) - - # Apply filter masking permanently back onto the base layer - target_conv.weight.copy_(original_weights * final_mask) - - # Unfreeze architecture parameters for evaluations downstream - for p in model.parameters(): - p.requires_grad = True - - print(f">> Permanently altered {out_channels} convolutional filters in layer4 via WF-Net.") - return model + diff --git a/unlearning/wf/WF_Net.py b/unlearning/wf/WF_Net.py index a5dd690..eb50180 100644 --- a/unlearning/wf/WF_Net.py +++ b/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