diff --git a/.gitignore b/.gitignore index 2a3ec8c..b100af8 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,5 @@ +# Created by venv; see https://docs.python.org/3/library/venv.html +* # Virtual Environment (the folders Git saw) bin/ lib/ @@ -19,4 +21,4 @@ lib64 *.idx *.rec *.lst -property \ No newline at end of file +property diff --git a/DataAnalyser.py b/DataAnalyser.py index 1c994a7..4dc4c2a 100644 --- a/DataAnalyser.py +++ b/DataAnalyser.py @@ -1,6 +1,6 @@ #from Data import * -from datasets.Casia import * +from sets.Casia import * ''' Because the size of samples per class had the biggest impact diff --git a/Tune.py b/Tune.py index 7f4d8e6..3abb32d 100644 --- a/Tune.py +++ b/Tune.py @@ -41,6 +41,7 @@ EPOCHS = 10 # Inception = 299 RESOLUTION = 224 +FINETUNE = False # whether to fintune or just load finetuned model from dir # model architecture options are # - RESNET18 # - RESNET50 @@ -112,19 +113,24 @@ device = SetUp.get_device() for i in range(0,1):#CLASS_SIZE): FORGET_CLASS_IDX = i # Create model using Factory - model = Model.create( - arch = arch, - device = device, - size = CLASS_SIZE) - # we may need to load existing model or finetune - #model.train( - # epochs = EPOCHS, - # loader = train_loader, - # rate = LR_RATE) + model = None - # save. - #model.save(filename=arch.name.lower()) + if FINETUNE: + model = Model.create( + arch = arch, + device = device, + size = CLASS_SIZE) + + # we may need to load existing model or finetune + model.train( + epochs = EPOCHS, + loader = train_loader, + rate = LR_RATE) + + # save. + file_name = f"{arch.name.lower}_{dataset_name.name.lower()}" + model.save(filename=arch.name.lower()) # done tuning @@ -147,18 +153,21 @@ for i in range(0,1):#CLASS_SIZE): # Evaluate current_mode = "Finetuned" - #accuracy, report_dict = model.evaluate( - # loader = test_loader, - # mode=current_mode - #) + if FINETUNE: - Util._log_to_csv( - arch=model.__class__.__name__, - mode = current_mode, - accuracy=accuracy, - report_dict=report_dict, - strategy="base" - ) + #current_mode = "Finetuned" + 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" + ) # unlearning algorithms #linear_filtration = LinearFiltration(target_class_index=FORGET_CLASS_IDX) @@ -167,7 +176,14 @@ for i in range(0,1):#CLASS_SIZE): #weight_filtration = WeightFiltration(num_classes = CLASS_SIZE,target_class_idx=FORGET_CLASS_IDX) #weight_filtration.apply(reloaded.model) - certified_removal = CertifiedRemoval(target_class_index=FORGET_CLASS_IDX,removal_bound=0.05, epsilon=0.5, l2_reg=15) + certified_removal = CertifiedRemoval( + target_class_index=FORGET_CLASS_IDX, + s1=2, + s2=500, + unlearn_bs=2, + scale=100.0, # Drop scale to match lower s2 depth + std=0.00001) + #,removal_bound=0.05, epsilon=0.5, l2_reg=15) #certified_removal.apply(reloaded.model) # to be unlearned @@ -200,6 +216,12 @@ for i in range(0,1):#CLASS_SIZE): # loader = test_loader #) + if not FINETUNE: + reloaded.evaluate( + loader = test_loader, + mode=current_mode + ) + # Unlearning # train loaders passed here strategy.apply(reloaded.model, forget_train_loader, retain_train_loader) diff --git a/Tune_new.py b/Tune_new.py index 1ca114a..d8f155d 100644 --- a/Tune_new.py +++ b/Tune_new.py @@ -10,6 +10,9 @@ 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 # Global Hyperparameters CLASS_SIZE = 20 @@ -17,7 +20,7 @@ BATCH_SIZE = 16 SAMPLE_SIZE = 30 TRAINING_SAMPLE = 27 RESOLUTION = 224 -ARCH = Architecture.RESNET50 +ARCH = Architecture.RESNET18 # Data preparation and model setup @@ -27,14 +30,17 @@ def prepare_data_and_model_environment(): train-test class splits, and configures the architecture base. """ device = SetUp.get_device() - dataset_name = Set_Name.CELEBA + dataset_name = Set_Name.CASIAFACES + if dataset_name == Set_Name.CASIAFACES: + SAMPLE_SIZE = 400 + TRAINING_SAMPLE = 320 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'> Selected {CLASS_SIZE} random identity classes from {dataset_name.name} dataset.') print(f'> A class has {TRAINING_SAMPLE} train and {SAMPLE_SIZE - TRAINING_SAMPLE} test samples') # Isolate sample index partitions @@ -81,6 +87,8 @@ def prepare_data_and_model_environment(): # 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. @@ -91,13 +99,15 @@ def run_finetuning_or_baseline_eval(env_dict, run_training=False, lr_rate=0.0001 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") + + train_loader = DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True) + + if not run_training: + return + #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)}") @@ -119,7 +129,7 @@ def run_finetuning_or_baseline_eval(env_dict, run_training=False, lr_rate=0.0001 # Unlearning and strategy eval -def run_unlearning_and_strategy_eval(env_dict, forget_class_idx): +def run_unlearning_and_strategy_eval(env_dict, forget_class_idx, strategy, evaluate = False): """ Reloads a clean model state, applies the isolated unlearning framework, and runs specific target evaluation domain checks. @@ -128,13 +138,34 @@ def run_unlearning_and_strategy_eval(env_dict, forget_class_idx): 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 - certified_removal = CertifiedRemoval( - target_class_index=forget_class_idx, - removal_bound=0.05, - epsilon=0.5, - l2_reg=15 - ) + # 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( @@ -147,11 +178,17 @@ def run_unlearning_and_strategy_eval(env_dict, forget_class_idx): # 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) + + if evaluate: + reloaded.evaluate( + loader=retain_test_loader, mode="finetuned" + ) + 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__ + strategy.apply(reloaded.model, forget_train_loader, retain_train_loader) + strategy_in_use = strategy.__class__.__name__ # Define validation tracking steps dynamically evaluation_domains = [ @@ -180,10 +217,62 @@ if __name__ == "__main__": 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=True) + 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 + # ) + + + + certified_unlearning = CertifiedUnlearning( + target_class_index=i, + l2_reg=0.000002, + gamma=0.1, + scale= 20000,# 16400.0, # took ages to reach this sweet spot + s1=2, + s2=300, + std=0.00001, + unlearn_bs=16 + ) + + # works perfectly + linear_filtration = LinearFiltration( + + target_class_index=i + ) + + weight_filtration = WeightFiltration( + target_class_index=i, + epochs=3, + lr=0.5, + gamma=150 + ) + + strategies = [ + certified_unlearning, + # weight_filtration, + # linear_filtration + ] + 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 + for strategy in strategies: + run_unlearning_and_strategy_eval( + runtime_environment, + forget_class_idx=i, + strategy=strategy, + evaluate= not finetuning + ) diff --git a/architectures/Model.py b/architectures/Model.py index 50e7771..287be1e 100644 --- a/architectures/Model.py +++ b/architectures/Model.py @@ -139,7 +139,7 @@ class Model(ABC): - # Using the factory patern here + # factory @staticmethod def create(arch, device, size): print(f'>> MODEL ARCHITECTURE >> {arch.name}.') @@ -151,6 +151,11 @@ class Model(ABC): from architectures.ResNet18 import ResNet18 return ResNet18(device, size) + # ResNet34 + case Architecture.RESNET34: + from architectures.ResNet34 import ResNet34 + return ResNet34(device, size) + # ResNet50 case Architecture.RESNET50: from architectures.ResNet50 import ResNet50 @@ -190,6 +195,7 @@ from enum import Enum, auto class Architecture(Enum): RESNET18 = auto() RESNET50 = auto() + RESNET34 = auto() INCEPTION = auto() DENSENET121 = auto() GOOGLENET = auto() diff --git a/sets/CelebA.py b/sets/CelebA.py index aff1f24..481dcc5 100644 --- a/sets/CelebA.py +++ b/sets/CelebA.py @@ -9,10 +9,10 @@ class CelebA(Data): def get_set(self): set = datasets.CelebA( - root = "./data", + root = "../data", split='all', target_type='identity', - download=True, + download=False, transform=None ) # set the target first diff --git a/sets/Extractor.py b/sets/Extractor.py index a3b2afd..dafe426 100644 --- a/sets/Extractor.py +++ b/sets/Extractor.py @@ -75,7 +75,7 @@ def extract_selected_binary(rec_path, idx_path, output_dir, top_labels): current_count = save_counters[label] img_filename = f"{current_count}.jpg" img_path = os.path.join(target_folder, img_filename) - if(current_count > 200): + if(current_count > 405): continue with open(img_path, 'wb') as img_f: @@ -119,9 +119,9 @@ if __name__ == "__main__": ''' if __name__ == "__main__": base_dir = os.path.dirname(os.path.abspath(__file__)) - REC = os.path.join(base_dir, 'casia', 'train.rec') - IDX = os.path.join(base_dir, 'casia', 'train.idx') - OUT = os.path.join(base_dir, 'casia-set') + REC = os.path.join(base_dir, '../data/casia-set', 'train.rec') + IDX = os.path.join(base_dir, '../data/casia-set', 'train.idx') + OUT = os.path.join(base_dir, '../data/casia-set') # Step 1: Trust the binary, not the text file top_verified_labels = get_top_identities_binary(REC, IDX, top_n=50) diff --git a/unlearning/CertifiedRemoval.py b/unlearning/CertifiedRemoval.py index 1395654..058fd22 100644 --- a/unlearning/CertifiedRemoval.py +++ b/unlearning/CertifiedRemoval.py @@ -1,127 +1,214 @@ import torch import torch.nn as nn -from torch.utils.data import DataLoader +from torch.utils.data import DataLoader, RandomSampler +from torch.autograd import grad from unlearning.Strategy import Strategy class CertifiedRemoval(Strategy): """ - Implements Certified Removal (Guo et al.) adapted for deep architectures - like ResNet50 by isolating and updating the final classification layer. + 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, removal_bound: float, epsilon: float, l2_reg: float = 0.1): - super().__init__() - self.removal_bound = removal_bound # gamma in the paper - self.epsilon = epsilon # Privacy budget - self.l2_reg = l2_reg # Lambda regularization term + 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 _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 = [] - all_labels = [] + ''' + 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 - with torch.no_grad(): - for inputs, labels in loader: - inputs = inputs.to(device) - # Pass through backbone to get the 2048-dimensional feature vector - features = backbone(inputs) - all_features.append(features.cpu()) - all_labels.append(labels.cpu()) + 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) - return torch.cat(all_features, dim=0), torch.cat(all_labels, dim=0) + 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: - """ - Entry point expected by your Model.unlearn() architecture interface. - Applies Certified Removal strictly to the final linear layer (model.fc). - """ + '''def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module: device = next(model.parameters()).device - - # Isolate the final NN (Fully connected) layer from the model - linear_head = model.fc - # Temporarily turn the fc layer into a identity pass-through - model.fc = nn.Identity() - - print(">> Extracting deep features from model backbone...") - retain_features, retain_labels = self._get_features(model, retain_loader, device) - forget_features, forget_labels = self._get_features(model, forget_loader, device) - - # Restore the linear head back - model.fc = linear_head - - # Extract weights from the classification layer - # w shape: [num_classes, 2048] - w = model.fc.weight.data.clone().cpu() - - # Compute the Exact Hessian Matrix over the remaining (retained) features - # Formula: H = (X^T * X) / N + lambda * I - N_retain = retain_features.size(0) - hessian = self._compute_hessian(retain_features=retain_features, retain_features_size = N_retain) - - grad_forget = self._compute_loss_gradient( - forget_labels=forget_labels, - forget_features=forget_features, - model_weights=w) - #torch.matmul(error.t(), forget_features) / forget_features.size(0) - # Compute the Newton step update via solving: H * Delta_W^T = Grad_forget^T - delta_w = self._compute_newton_step( - tensor = hessian, - gradient= grad_forget - ) - # Apply the Certified Removal update rule: W_new = W + Delta_W - new_w = w + delta_w - # Calibrate noise based on your epsilon budget - # (Guo et al. use a perturbation based on the regularization lambda and epsilon) - sigma = 2.0 / (self.l2_reg * self.epsilon) - noise = torch.randn_like(new_w) * (sigma / N_retain) - new_w = new_w + noise - - # Theoretical Guarantee verification - 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. Norm: {norm_delta:.4f} <= Bound: {self.removal_bound}") - - # Push updated parameters back into the model instance in-place - model.fc.weight.data = new_w.to(device) + num_forget = len(forget_loader.dataset) + num_retain = len(retain_loader.dataset) + scaling_ratio = num_forget / num_retain - print(">> Certified Removal process completed successfully.") - return model - - - # computing the hessian matrix - def _compute_hessian(self, retain_features, retain_features_size): - 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)) - return (X_T_X / retain_features_size) + reg_matrix - - - def _compute_loss_gradient(self, forget_features, forget_labels, model_weights): - print(">> Calculating forget set gradients...") - num_classes = model_weights.size(0) - # Pass features through linear layer weights to get logits - logits_forget = torch.matmul(forget_features, model_weights.t()) - # Apply softmax to get true class probabilities - preds_softmax = torch.softmax(logits_forget, dim=1) + 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) - forget_labels_one_hot = torch.nn.functional.one_hot(forget_labels, num_classes=num_classes).float() + 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 - error = preds_softmax - forget_labels_one_hot - # grad_forget shape: [num_classes, 2048] - return torch.matmul(error.t(), forget_features) / forget_features.size(0) - - - def _compute_newton_step(self,tensor, gradient): - print(">> Solving Newton step via system optimization...") - try: - delta_w_t = torch.linalg.solve(tensor, gradient.t()) - delta_w = delta_w_t.t() - except RuntimeError: - print(">> Warning: Hessian matrix is singular. Falling back to pseudo-inverse.") - delta_w = torch.matmul(gradient, torch.linalg.pinv(tensor).t()) - return delta_w \ No newline at end of file + 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/LinearFiltration.py b/unlearning/LinearFiltration.py index 3218c3e..ebf9343 100644 --- a/unlearning/LinearFiltration.py +++ b/unlearning/LinearFiltration.py @@ -1,47 +1,184 @@ - import torch import torch.nn as nn from .Strategy import Strategy from torch.utils.data import DataLoader class LinearFiltration(Strategy): - def __init__(self,target_class_index): - super().__init__(target_class_index = target_class_index) + 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() + # Freeze internal params for param in model.parameters(): param.requires_grad = False + + device = next(model.parameters()).device - with torch.no_grad(): - W = model.fc.weight.data.clone() - num_classes = W.shape[0] - - 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) - if model.fc.bias is not None: - b = model.fc.bias.data.clone() - # b is a 1D tensor of shape (num_classes), - # so we use torch.mv (matrix-vector multiplication) or unsqueeze it - sanitized_b = torch.mv(A, b) - model.fc.bias.copy_(sanitized_b) - - return model + return self.normalise( + model=model, + retain_loader=retain_loader, + forget_loader=forget_loader, + device=device, + 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 - # The row of the forgotten class should average all other classes for j in range(num_classes): if j == forget_class: - # we zero the forget class A[forget_class, j] = 0.0 else: - # and we distribute the output to the remaining A[forget_class, j] = 1.0 / num_remaining - return A \ No newline at end of file + return A + + + @staticmethod + def _sums_and_counts(model, num_classes, retain_loader, forget_loader, device, forget_index, h_dim): + model.eval() + + sums = torch.zeros(num_classes, h_dim, device=device) + counts = torch.zeros(num_classes, device=device) + + # Generate values for retain + with torch.no_grad(): + for inputs, targets in retain_loader: + inputs = inputs.to(device) + targets = targets.to(device) + # FIX: Call get_features instead of model() directly + outputs = LinearFiltration.get_features(model, inputs) + + for j in range(num_classes): + if j == forget_index: + continue + mask = (targets == j) + + if mask.any(): + 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( + model=model, + num_classes=num_classes, + retain_loader=retain_loader, + forget_loader=forget_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) + + + @staticmethod + def _compute_z(tensor, forget_index): + # Now tensor has shape (num_classes, h_dim) -> tensor.shape[0] is num_classes + K = tensor.shape[0] + + # pi_a0 should match the feature space dimensions (h_dim) + pi_a0 = torch.zeros(tensor.shape[1], device=tensor.device) + + t_1 = pi_a0 + a0 = tensor[forget_index, :] # Extracting the row vector for the forgotten class + + mask_a0 = torch.ones( + a0.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() + + mask_rows = torch.ones(K, dtype=torch.bool, device=tensor.device) + mask_rows[forget_index] = False + + r_A = tensor[mask_rows, :] + t_3 = (1.0 / ((K - 1)) ** 2) * r_A.sum() + + return t_1 + t_2 + t_3 + + + @staticmethod + def normalise(model, retain_loader, forget_loader, device, forget_index): + W = model.fc.weight.data.clone() + num_classes = W.shape[0] + + A = LinearFiltration._get_means( + model=model, + num_classes=num_classes, + retain_loader=retain_loader, + forget_loader=forget_loader, + device=device, + forget_index=forget_index + ) + + Z = LinearFiltration._compute_z(tensor=A, forget_index=forget_index) + B_Z_rows = [] + + for i in range(num_classes): + if i == forget_index: + B_Z_rows.append(Z) + else: + # Retained classes maintain their original ideal feature directions + B_Z_rows.append(A[i]) + + # Stack back along dim=0 to match (num_classes, h_dim) + B_Z = torch.stack(B_Z_rows, dim=0) + + A_inv = torch.linalg.pinv(A) + + W_Z = B_Z @ A_inv @ W + + model.fc.weight.copy_(W_Z) + + return model \ No newline at end of file diff --git a/unlearning/WeightFiltration.py b/unlearning/WeightFiltration.py index c806020..3dd03eb 100644 --- a/unlearning/WeightFiltration.py +++ b/unlearning/WeightFiltration.py @@ -3,97 +3,34 @@ import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader from unlearning.Strategy import Strategy +from .wf.WF_Net import WF_Net 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. + 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,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 + def __init__(self, target_class_index: int, epochs: int = 10, lr: float = 0.2, gamma: float = 10.0): + super().__init__(target_class_index=target_class_index) self.epochs = epochs self.lr = lr self.gamma = gamma - self.alpha = None - self.hook_handle = None + #self.alpha = None - - def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module: - device = next(model.parameters()).device - model.eval() + + + + + 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) - # Locate layer4 for dynamic optimization - target_layer = model.layer4 if hasattr(model, 'layer4') else None - fc_layer = model.fc if hasattr(model, 'fc') and isinstance(model.fc, nn.Linear) else None - - if target_layer is None or fc_layer is None: - raise AttributeError("Model does not have the required layers.") - - # Match alpha dimensions to the channels outputted by layer4 - num_features = fc_layer.weight.shape[1] - self.alpha = nn.Parameter(torch.ones(self.num_classes, num_features, device=device) * 1.5) - - # Freeze everything except our channel mask - for p in model.parameters(): - p.requires_grad = False - self.alpha.requires_grad = True - - # Hook into layer4 dynamically to run the untraining optimization - self.hook_handle = target_layer.register_forward_hook(self._get_hook()) - # optimise the filter to maintain accuracy on retain set - # and decrease accuracy on forget set - self._optimise_filter(model, forget_loader, retain_loader, device) - - # Remove the runtime hook - self.hook_handle.remove() - - # Transfer the channel suppression permanently into model.fc - with torch.no_grad(): - 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_index].copy_( - fc_layer.weight[self.target_class_index] * mask - ) - print(f">> Baked deep channel filter into Class {self.target_class_index} weights.") - - return model - - def _get_hook(self): - """ - Filters the internal feature map channels of layer4. - The mask scales the channels across the batch. - """ - def functional_hook(module, layer_input, layer_output): - # layer_output shape: (batch, channels, height, width) -> e.g., (16, 2048, 7, 7) - # 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_index]) - - # Reshape mask to (1, channels, 1, 1) so it broadcasts over batch, height, and width - mask = mask.view(1, -1, 1, 1) - - # Scale the internal feature maps before they move to the next layer - return layer_output * mask - - return functional_hook - - - def _optimise_filter(self, model, forget_loader, retain_loader, device): - optimizer = optim.Adam([self.alpha], lr=self.lr) + # 2. ONLY optimize alpha (everything else is frozen inside the wrapper) + optimizer = optim.Adam([wf_model.alpha], lr=self.lr) criterion = nn.CrossEntropyLoss() - 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 - # this can be adjusted to optimise the best escape point - # it is the value we set to evaluate performance improvement after each itteration. - # if improvement is less than this, then we break itteration. - threshold = 0.05 - for epoch in range(self.epochs): forget_iter = iter(forget_loader) t_loss_r, t_loss_f = 0.0, 0.0 @@ -102,6 +39,7 @@ class WeightFiltration(Strategy): for r_inputs, r_labels in retain_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: @@ -111,10 +49,19 @@ class WeightFiltration(Strategy): optimizer.zero_grad() - # Compute Losses - # The hook handles the weight filtering smoothly behind the scenes - loss_r = criterion(model(r_inputs), r_labels) - loss_f = -torch.sum((torch.ones_like(model(f_inputs)) / self.num_classes) * torch.log_softmax(model(f_inputs), dim=-1)) + # --- 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) + + # Compute Losses using Poppi et al.'s temperature scaled entropy + loss_r = criterion(outputs_r, r_labels) + + temperature = 3.0 + logits_f_scaled = outputs_f / temperature + loss_f = -torch.sum( + (torch.ones_like(logits_f_scaled) / num_classes) * torch.log_softmax(logits_f_scaled, dim=-1) + ) total_loss = loss_r + (self.gamma * loss_f) total_loss.backward() @@ -122,17 +69,56 @@ class WeightFiltration(Strategy): t_loss_r += loss_r.item() t_loss_f += loss_f.item() - steps += 1 - forget_loss = t_loss_f / steps - print(f" Epoch {epoch+1}/{self.epochs} | Retain Loss: {t_loss_r/steps:.4f} | Forget Loss: {forget_loss:.4f}") + + print(f" Epoch {epoch+1}/{self.epochs} | Retain Loss: {t_loss_r/steps:.4f} | Forget Loss: {t_loss_f/steps:.4f}") - if temp_forget_loss is not None: + return wf_model + - improvement = temp_forget_loss - forget_loss - # if optimisation reaches a point of diminishing returns (improvements is less than threshold) - # we break the loop - if improvement < threshold: - break - # else we update the lasst recorded loss. - temp_forget_loss = forget_loss \ No newline at end of file + + def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module: + device = next(model.parameters()).device + 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 + else: + raise AttributeError("Model architecture does not match expected ResNet-18 structure.") + + # 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, + ) + + # --- PERMANENT BAKING STEP --- + # Disconnect the dynamic parameter and freeze the optimal gated state permanently into the architecture + + with torch.no_grad(): + final_mask = torch.sigmoid(wf_model.alpha[self.target_class_index]).view(-1, 1, 1, 1) + target_conv.weight.copy_(original_weights * final_mask) + + # Re-enable model parameters for downstream evaluation processing + for p in model.parameters(): + p.requires_grad = True + + print(f">> Permanently altered {out_channels} convolutional filters in layer4 via WF-Net.") + return model \ No newline at end of file