From 207fcae6991d16cc9b1f8f3077d618db2e237746 Mon Sep 17 00:00:00 2001 From: Tinsae Date: Sun, 14 Jun 2026 23:13:33 +0200 Subject: [PATCH] facebook's implementation --- Tune.py | 40 +++---- Tune_new.py | 189 +++++++++++++++++++++++++++++++ unlearning/CertifiedRemoval.py | 26 ----- unlearning/Certified_facebook.py | 123 ++++++++++++++++++++ unlearning/LinearFiltration.py | 7 +- unlearning/Strategy.py | 3 +- unlearning/WeightFiltration.py | 18 ++- 7 files changed, 345 insertions(+), 61 deletions(-) create mode 100644 Tune_new.py create mode 100644 unlearning/Certified_facebook.py diff --git a/Tune.py b/Tune.py index 3fb1b6f..7f4d8e6 100644 --- a/Tune.py +++ b/Tune.py @@ -34,7 +34,7 @@ TRAINING_SMPLE = 27 # learning rate LR_RATE = 0.0001 -EPOCHS = 20 +EPOCHS = 10 # depends on model architecture # ResNet, DenseNet = 224 @@ -109,7 +109,7 @@ print(f'> Constants : Classes = {CLASS_SIZE}, Batch = {BATCH_SIZE}, epochs = {EP device = SetUp.get_device() -for i in range(0,CLASS_SIZE): +for i in range(0,1):#CLASS_SIZE): FORGET_CLASS_IDX = i # Create model using Factory model = Model.create( @@ -118,13 +118,13 @@ for i in range(0,CLASS_SIZE): size = CLASS_SIZE) # we may need to load existing model or finetune - model.train( - epochs = EPOCHS, - loader = train_loader, - rate = LR_RATE) + #model.train( + # epochs = EPOCHS, + # loader = train_loader, + # rate = LR_RATE) # save. - model.save(filename=arch.name.lower()) + #model.save(filename=arch.name.lower()) # done tuning @@ -147,10 +147,10 @@ for i in range(0,CLASS_SIZE): # Evaluate current_mode = "Finetuned" - accuracy, report_dict = model.evaluate( - loader = test_loader, - mode=current_mode - ) + #accuracy, report_dict = model.evaluate( + # loader = test_loader, + # mode=current_mode + #) Util._log_to_csv( arch=model.__class__.__name__, @@ -161,13 +161,13 @@ for i in range(0,CLASS_SIZE): ) # unlearning algorithms - linear_filtration = LinearFiltration(target_class_idx=FORGET_CLASS_IDX) + #linear_filtration = LinearFiltration(target_class_index=FORGET_CLASS_IDX) #filtration.apply(reloaded.model) - weight_filtration = WeightFiltration(num_classes = CLASS_SIZE,target_class_idx=FORGET_CLASS_IDX) + #weight_filtration = WeightFiltration(num_classes = CLASS_SIZE,target_class_idx=FORGET_CLASS_IDX) #weight_filtration.apply(reloaded.model) - certified_removal = CertifiedRemoval(removal_bound=0.05, epsilon=0.5, l2_reg=0.1) + certified_removal = CertifiedRemoval(target_class_index=FORGET_CLASS_IDX,removal_bound=0.05, epsilon=0.5, l2_reg=15) #certified_removal.apply(reloaded.model) # to be unlearned @@ -179,14 +179,14 @@ for i in range(0,CLASS_SIZE): # to evaluate forget_test_loader, retain_test_loader = get_unlearning_loaders( - dataset=test_data, - forget_class_idx=FORGET_CLASS_IDX, - batch_size=BATCH_SIZE - ) + dataset=test_data, + forget_class_idx=FORGET_CLASS_IDX, + batch_size=BATCH_SIZE + ) - strategies = [linear_filtration, weight_filtration, certified_removal] - #strategies = [linear_filtration] + #strategies = [linear_filtration, weight_filtration, certified_removal] + strategies = [certified_removal] for strategy in strategies: # test again reloaded = Model.create( diff --git a/Tune_new.py b/Tune_new.py new file mode 100644 index 0000000..1ca114a --- /dev/null +++ b/Tune_new.py @@ -0,0 +1,189 @@ +import torch +import torch.nn as nn +from torch.utils.data import DataLoader +from sklearn.metrics import classification_report + +# Framework and Utility Imports +import SetUp +import Util +from sets.Data import * +from sets.IdentitySubset import IdentitySubset +from architectures.Model import Model, Architecture +from unlearning.CertifiedRemoval import CertifiedRemoval + +# Global Hyperparameters +CLASS_SIZE = 20 +BATCH_SIZE = 16 +SAMPLE_SIZE = 30 +TRAINING_SAMPLE = 27 +RESOLUTION = 224 +ARCH = Architecture.RESNET50 + + +# Data preparation and model setup +def prepare_data_and_model_environment(): + """ + Handles environment discovery, downloads/loads datasets, generates + train-test class splits, and configures the architecture base. + """ + device = SetUp.get_device() + dataset_name = Set_Name.CELEBA + + dataset = get_set(set_name=dataset_name) + print(f"> {dataset.__class__.__name__} dataset loaded") + + # Select target identities (deterministic top sample identities) + selected_identities = select_top_ids(dataset=dataset, class_size=CLASS_SIZE) + print(f'> Selected {CLASS_SIZE} random identity classes from CelebA dataset.') + print(f'> A class has {TRAINING_SAMPLE} train and {SAMPLE_SIZE - TRAINING_SAMPLE} test samples') + + # Isolate sample index partitions + train_indices, test_indices = get_indices( + dataset=dataset, + identities=selected_identities, + split_at=TRAINING_SAMPLE, + size=SAMPLE_SIZE + ) + + # Remap identities to 0 -> (N-1) range required by CrossEntropyLoss + id_map = {old_id: new_id for new_id, old_id in enumerate(selected_identities)} + + # Build internal datasets using custom transforms + tr_transform = train_transform(RESOLUTION) + train_data = IdentitySubset( + dataset=dataset, + indices=train_indices, + id_mapping=id_map, + transform=tr_transform + ) + + te_transform = test_transform(RESOLUTION) + test_data = IdentitySubset( + dataset=dataset, + indices=test_indices, + id_mapping=id_map, + transform=te_transform + ) + + print(f"> Total training images: {len(train_data)}") + print(f'> Constants : Classes = {CLASS_SIZE}, Batch = {BATCH_SIZE}') + + # Create the base target model instance + base_model = Model.create(arch=ARCH, device=device, size=CLASS_SIZE) + + return { + "device": device, + "train_data": train_data, + "test_data": test_data, + "base_model": base_model + } + + +# Fine tunning and evaluation +def run_finetuning_or_baseline_eval(env_dict, run_training=False, lr_rate=0.0001, epochs=10): + """ + Handles model training (if flag is true) and logs the baseline fine-tuned + performance to file metrics. + """ + model = env_dict["base_model"] + train_data = env_dict["train_data"] + test_data = env_dict["test_data"] + + test_loader = DataLoader(test_data, batch_size=BATCH_SIZE, shuffle=False) + + # Optional training configuration switch + if run_training: + train_loader = DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True) + print(f"Starting training on {env_dict['device']}...") + model.train(epochs=epochs, loader=train_loader, rate=lr_rate) + model.save(filename=ARCH.name.lower()) + print(f"Model saved to trained_models/{ARCH.name.lower()}.pth") + + print(f"Total test images for these {CLASS_SIZE} classes: {len(test_data)}") + + # Evaluate original base checkpoint performance + current_mode = "Finetuned" + + # Check if weights exist or model was trained before evaluating + try: + accuracy, report_dict = model.evaluate(loader=test_loader, mode=current_mode) + Util._log_to_csv( + arch=model.__class__.__name__, + mode=current_mode, + accuracy=accuracy, + report_dict=report_dict, + strategy="base" + ) + except Exception as e: + print(f">> Skipping baseline log generation. Reason: {e}") + + +# Unlearning and strategy eval +def run_unlearning_and_strategy_eval(env_dict, forget_class_idx): + """ + Reloads a clean model state, applies the isolated unlearning framework, + and runs specific target evaluation domain checks. + """ + device = env_dict["device"] + train_data = env_dict["train_data"] + test_data = env_dict["test_data"] + + # Initialize the strategy hyperparameters matching standard settings + certified_removal = CertifiedRemoval( + target_class_index=forget_class_idx, + removal_bound=0.05, + epsilon=0.5, + l2_reg=15 + ) + + # Segment specific unlearning loaders using class index boundaries + forget_train_loader, retain_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( + dataset=test_data, forget_class_idx=forget_class_idx, batch_size=BATCH_SIZE + ) + + # Instantiate a clean copy of the model to keep weights isolated + reloaded = Model.create(arch=ARCH, device=device, size=CLASS_SIZE) + reloaded.load(arch=ARCH) + print("fine tunned model loaded into evaluation sandbox") + + # Execute strategic parameter unlearning step + certified_removal.apply(reloaded.model, forget_train_loader, retain_train_loader) + strategy_in_use = certified_removal.__class__.__name__ + + # Define validation tracking steps dynamically + evaluation_domains = [ + {"loader": retain_test_loader, "mode": "retain", "label": "\n--- Performance on Retained Classes"}, + {"loader": forget_test_loader, "mode": "forget", "label": "\n--- Performance on Forgotten Class"}, + {"loader": forget_train_loader, "mode": "forget_train", "label": "\n--- Performance on Forgotten Class (Train Set - Verifying Unlearning)"} + ] + + # Process and append metrics to target reporting paths + for domain in evaluation_domains: + print(domain["label"]) + accuracy, report_dict = reloaded.evaluate(loader=domain["loader"], mode=domain["mode"]) + Util._log_to_csv( + arch=reloaded.__class__.__name__, + mode=domain["mode"], + accuracy=accuracy, + report_dict=report_dict, + strategy=strategy_in_use + ) + + +# entry +if __name__ == "__main__": + + # Run Data Infrastructure and Architecture Builder + runtime_environment = prepare_data_and_model_environment() + + # Baseline Evaluation + # switch finetuning for tests on strategies only + run_finetuning_or_baseline_eval(runtime_environment, run_training=True) + + # Unlearning Iterations + for i in range(0, 1): + print(f"\n>>> Executing Unlearning Framework for Target Identity Index: {i} <<<") + run_unlearning_and_strategy_eval(runtime_environment, forget_class_idx=i) \ No newline at end of file diff --git a/unlearning/CertifiedRemoval.py b/unlearning/CertifiedRemoval.py index 8ef060a..1395654 100644 --- a/unlearning/CertifiedRemoval.py +++ b/unlearning/CertifiedRemoval.py @@ -55,27 +55,9 @@ class CertifiedRemoval(Strategy): # Compute the Exact Hessian Matrix over the remaining (retained) features # Formula: H = (X^T * X) / N + lambda * I - # this will be done on CPU. requires more ram so we cant afford to do it on VRAM - # print(">> Computing exact Hessian matrix...") N_retain = retain_features.size(0) - # X_T_X = torch.matmul(retain_features.t(), retain_features) - # reg_matrix = self.l2_reg * torch.eye(retain_features.size(1)) hessian = self._compute_hessian(retain_features=retain_features, retain_features_size = N_retain) - # Compute the gradient of the loss with respect to the forgotten data - # print(">> Calculating forget set gradients...") - # num_classes = w.size(0) - # Pass features through linear layer weights to get logits - # logits_forget = torch.matmul(forget_features, w.t()) - # Apply softmax to get true class probabilities - # preds_softmax = torch.softmax(logits_forget, dim=1) - - # forget_labels_one_hot = torch.nn.functional.one_hot(forget_labels, num_classes=num_classes).float() - - #preds_forget = torch.matmul(forget_features, w.t()) - #error = preds_forget - forget_labels_one_hot - # error = preds_softmax - forget_labels_one_hot - # grad_forget shape: [num_classes, 2048] grad_forget = self._compute_loss_gradient( forget_labels=forget_labels, forget_features=forget_features, @@ -87,14 +69,6 @@ class CertifiedRemoval(Strategy): tensor = hessian, gradient= grad_forget ) - # print(">> Solving Newton step via system optimization...") - # try: - # delta_w_t = torch.linalg.solve(Hessian, grad_forget.t()) - # delta_w = delta_w_t.t() - # except RuntimeError: - # print(">> Warning: Hessian matrix is singular. Falling back to pseudo-inverse.") - # delta_w = torch.matmul(grad_forget, torch.linalg.pinv(Hessian).t()) - # Apply the Certified Removal update rule: W_new = W + Delta_W new_w = w + delta_w # Calibrate noise based on your epsilon budget diff --git a/unlearning/Certified_facebook.py b/unlearning/Certified_facebook.py new file mode 100644 index 0000000..0b8bcd1 --- /dev/null +++ b/unlearning/Certified_facebook.py @@ -0,0 +1,123 @@ +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/LinearFiltration.py b/unlearning/LinearFiltration.py index 6bd8de2..3218c3e 100644 --- a/unlearning/LinearFiltration.py +++ b/unlearning/LinearFiltration.py @@ -5,9 +5,8 @@ from .Strategy import Strategy from torch.utils.data import DataLoader class LinearFiltration(Strategy): - def __init__(self, target_class_idx: int): - super().__init__() - self.target_class_idx = target_class_idx + def __init__(self,target_class_index): + super().__init__(target_class_index = target_class_index) def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module: model.eval() @@ -18,7 +17,7 @@ class LinearFiltration(Strategy): W = model.fc.weight.data.clone() num_classes = W.shape[0] - A = self._calculate_filtration_matrix(num_classes, self.target_class_idx, W.device) + A = self._calculate_filtration_matrix(num_classes, self.target_class_index, W.device) sanitized_W = torch.mm(A, W) model.fc.weight.copy_(sanitized_W) # Filter the bias (if the layer uses one) diff --git a/unlearning/Strategy.py b/unlearning/Strategy.py index e96a932..1e29255 100644 --- a/unlearning/Strategy.py +++ b/unlearning/Strategy.py @@ -9,9 +9,10 @@ import Util class Strategy: """Abstract base class for unlearning algorithms with automated, strategy-specific logging.""" - def __init__(self): + def __init__(self, target_class_index): # Dynamically set file name based on the class name (e.g., 'NormalizingLinearFiltration.txt') self.strategy_name = self.__class__.__name__ + self.target_class_index = target_class_index self.log_file = Path(f"reports/{self.strategy_name}/metrics.txt") Util._initialize_log_file(log_file= self.log_file) diff --git a/unlearning/WeightFiltration.py b/unlearning/WeightFiltration.py index ecf49da..c806020 100644 --- a/unlearning/WeightFiltration.py +++ b/unlearning/WeightFiltration.py @@ -1,6 +1,5 @@ import torch import torch.nn as nn -import torch.nn.functional as F import torch.optim as optim from torch.utils.data import DataLoader from unlearning.Strategy import Strategy @@ -10,10 +9,9 @@ class WeightFiltration(Strategy): Implements Poppi et al.'s Weight Filtering framework for linear layers. Uses a standard functional hook to guarantee native PyTorch autograd tracking. """ - def __init__(self, num_classes: int, target_class_idx: int, epochs: int = 10, lr: float = 0.2, gamma: float = 10.0): - super().__init__() + def __init__(self, target_class_index,num_classes: int, epochs: int = 10, lr: float = 0.2, gamma: float = 10.0): + super().__init__(target_class_index = target_class_index) self.num_classes = num_classes - self.target_class_idx = target_class_idx self.epochs = epochs self.lr = lr self.gamma = gamma @@ -52,13 +50,13 @@ class WeightFiltration(Strategy): # Transfer the channel suppression permanently into model.fc with torch.no_grad(): - mask = torch.sigmoid(self.alpha[self.target_class_idx]) # Shape: (num_features,) + mask = torch.sigmoid(self.alpha[self.target_class_index]) # Shape: (num_features,) # Suppress the channels ONLY for the target class row in fc - fc_layer.weight[self.target_class_idx].copy_( - fc_layer.weight[self.target_class_idx] * mask + fc_layer.weight[self.target_class_index].copy_( + fc_layer.weight[self.target_class_index] * mask ) - print(f">> Baked deep channel filter into Class {self.target_class_idx} weights.") + print(f">> Baked deep channel filter into Class {self.target_class_index} weights.") return model @@ -72,7 +70,7 @@ class WeightFiltration(Strategy): # self.alpha shape: (num_classes, channels) -> e.g., (20, 2048) # Extract 1D mask for the target class: (channels,) - mask = torch.sigmoid(self.alpha[self.target_class_idx]) + mask = torch.sigmoid(self.alpha[self.target_class_index]) # Reshape mask to (1, channels, 1, 1) so it broadcasts over batch, height, and width mask = mask.view(1, -1, 1, 1) @@ -87,7 +85,7 @@ class WeightFiltration(Strategy): optimizer = optim.Adam([self.alpha], lr=self.lr) criterion = nn.CrossEntropyLoss() - print(f"[{self.__class__.__name__}] Unlearning Class {self.target_class_idx} with gamma={self.gamma}...") + print(f"[{self.__class__.__name__}] Unlearning Class {self.target_class_index} with gamma={self.gamma}...") # To optimise this loop we will watch improvements after each optimisation temp_forget_loss = None