From 434d3b8198569b7f6791d429dcb62f2879f6a6b8 Mon Sep 17 00:00:00 2001 From: Tinsae Date: Wed, 1 Jul 2026 21:05:01 +0200 Subject: [PATCH] optimised --- .gitignore | 24 +++ DataAnalyser.py | 50 +++++ Note.md | 40 ++++ Predict.py | 67 ++++++ ReadME.md | 85 ++++++++ SetUp.py | 14 ++ Tune.py | 268 +++++++++++++++++++++++ Tune_new.py | 285 +++++++++++++++++++++++++ Util.py | 51 +++++ architectures/DenseNet121.py | 15 ++ architectures/EfficentNet.py | 19 ++ architectures/GoogleNet.py | 31 +++ architectures/Inception.py | 47 ++++ architectures/Model.py | 225 +++++++++++++++++++ architectures/ResNet18.py | 22 ++ architectures/ResNet34.py | 22 ++ architectures/ResNet50.py | 36 ++++ architectures/ShuffleNet.py | 17 ++ architectures/WFNet.py | 230 ++++++++++++++++++++ architectures/WideResNet.py | 15 ++ dependencies.txt | 6 + js_evaluator/JS_Evaluator.py | 111 ++++++++++ sets/Casia.py | 167 +++++++++++++++ sets/CasiaFace.py | 21 ++ sets/CelebA.py | 20 ++ sets/Data.py | 239 +++++++++++++++++++++ sets/Data_OOP.py | 174 +++++++++++++++ sets/Extractor.py | 131 ++++++++++++ sets/IdentitySubset.py | 34 +++ unlearning/CertifiedUnlearning.py | 344 ++++++++++++++++++++++++++++++ unlearning/LinearFiltration.py | 168 +++++++++++++++ unlearning/Strategy.py | 62 ++++++ unlearning/WF.py | 107 ++++++++++ unlearning/WeightFiltration.py | 139 ++++++++++++ 34 files changed, 3286 insertions(+) create mode 100644 .gitignore create mode 100644 DataAnalyser.py create mode 100644 Note.md create mode 100644 Predict.py create mode 100644 ReadME.md create mode 100644 SetUp.py create mode 100644 Tune.py create mode 100644 Tune_new.py create mode 100644 Util.py create mode 100644 architectures/DenseNet121.py create mode 100644 architectures/EfficentNet.py create mode 100644 architectures/GoogleNet.py create mode 100644 architectures/Inception.py create mode 100644 architectures/Model.py create mode 100644 architectures/ResNet18.py create mode 100644 architectures/ResNet34.py create mode 100644 architectures/ResNet50.py create mode 100644 architectures/ShuffleNet.py create mode 100644 architectures/WFNet.py create mode 100644 architectures/WideResNet.py create mode 100644 dependencies.txt create mode 100644 js_evaluator/JS_Evaluator.py create mode 100644 sets/Casia.py create mode 100644 sets/CasiaFace.py create mode 100644 sets/CelebA.py create mode 100644 sets/Data.py create mode 100644 sets/Data_OOP.py create mode 100644 sets/Extractor.py create mode 100644 sets/IdentitySubset.py create mode 100644 unlearning/CertifiedUnlearning.py create mode 100644 unlearning/LinearFiltration.py create mode 100644 unlearning/Strategy.py create mode 100644 unlearning/WF.py create mode 100644 unlearning/WeightFiltration.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..39815d0 --- /dev/null +++ b/.gitignore @@ -0,0 +1,24 @@ +# Created by venv; see https://docs.python.org/3/library/venv.html +#* +# Virtual Environment (the folders Git saw) +bin/ +lib/ +share/ +pyvenv.cfg +include/ + +# Data and Models +data/ +datasets/ +trained_models/ + +# Python cache +__pycache__/ +*.py[cod] +lib64 +/reports +*.bin +*.idx +*.rec +*.lst +property diff --git a/DataAnalyser.py b/DataAnalyser.py new file mode 100644 index 0000000..4dc4c2a --- /dev/null +++ b/DataAnalyser.py @@ -0,0 +1,50 @@ + +#from Data import * +from sets.Casia import * + +''' +Because the size of samples per class had the biggest impact +on training outcome, I decided to check the maximum amount of data +I can get from a class. +The highest I can get is +Rank | Identity ID |Count +----------------------------------- +1 | 3782 | 35 +2 | 2820 | 35 +3 | 3227 | 35 +4 | 3745 | 34 +5 | 3699 | 34 +6 | 8968 | 32 +7 | 9152 | 32 +8 | 9256 | 32 +9 | 2114 | 31 +... | ... | ... +17 | 4126 | 31 +18 | 3185 | 30 +... | ... | ... +50 | 3186 | 30 + +as can be seen, 3 classes have 35, 2 have 34, 3 have 32 and the rest have 30. +''' +def print_top_identity_stats(dataset, top_n=50): + # we get data + ids, counts = get_ids_and_counts(dataset) + # sort in descending order + sorted_counts, sorted_indices = torch.sort(counts, descending=True) + + # coresponding sorted ids + sorted_ids = ids[sorted_indices] + + # 4. Slice the first 'top_n' and print + print(f"{'Rank':<8} | {'Identity ID':<12} | {'Image Count':<12}") + print("-" * 35) + + for i in range(top_n): + identity_id = sorted_ids[i].item() + count = sorted_counts[i].item() + print(f"{i+1:<8} | {identity_id:<12} | {count:<12}") + +# Usage: +dataset = get_set() +print_top_identity_stats(dataset, 50) + diff --git a/Note.md b/Note.md new file mode 100644 index 0000000..67e3557 --- /dev/null +++ b/Note.md @@ -0,0 +1,40 @@ +Not at all! You are still completely faithful to Guo et al. Your current implementation does **not** break their framework. + +In fact, your setup matches the exact methodology of the paper. There is a common misconception about what Guo et al. mean when they caution against calculating the Hessian inverse, and understanding how your feature extraction relates to their theory explains why your code remains completely valid. + +--- + +## 1. What Guo et al. *Actually* Said + +In the original paper (*"Certified Removal from Linear Models"*), Guo et al. state that explicitly calculating and inverting the Hessian matrix becomes prohibitively expensive when the parameter count $d$ scales up. + +$$\text{Time Complexity to Invert } H = O(d^3)$$ + +However, the authors explicitly implemented and verified their approach on **linear classifiers** (like logistic regression) where the input feature dimension $d$ was small enough to handle directly. + +When you strip out the heavy ResNet50 convolutional layers and turn the backbone into a static feature extractor, **you transform your deep network into a linear classifier.** ``` +[Images] ──> [Frozen ResNet Backbone] ──> Extracted Feature Vector (d = 2048) ──> [Linear Head (model.fc)] + + +Because your feature vector is exactly 2,048 dimensions, your Hessian matrix is a modest $2048 \times 2048$. + +Inverting a $2048 \times 2048$ matrix takes your CPU less than **0.5 seconds** ($2048^3$ operations is tiny for a modern processor). You are executing the exact mathematics Guo et al. prescribed for linear systems. You haven't broken their implementation; you have successfully reduced a non-convex deep learning problem into their exact convex linear domain. + + +## 2. Where Hessian-Free Approximations (Like LiSSA) Apply + +The reason alternative methods like LiSSA or Conjugate Gradient exist is for scenarios where you *cannot* reduce the model to a small linear head. + +If you decided to apply Certified Removal to the **entire ResNet50 network** (all 23.5 million parameters open), then you would be forced to abandon your exact matrix calculation. Inverting a $23.5\text{M} \times 23.5\text{M}$ matrix is impossible. In that specific scenario, you would have to use a Hessian-free approximation method to avoid breaking the budget. + +--- + +## 3. The Core Alignment with the Paper + +Your script implements the three pillars that define Guo et al.’s Certified Removal: + +1. **The Optimization Target:** It uses an $L_2$-regularized objective function (`self.l2_reg`). +2. **The Newton Step:** It takes the exact second-order curvature correction ($H^{-1} \nabla$) to adjust the parameters. +3. **The Indistinguishability Guarantee:** It applies a privacy perturbation boundary check (`self.removal_bound`). + +Your implementation is an elegant, academically sound adaptation of their linear model theory for a deep learning architecture. By handling the feature extraction step first, you made their exact algorithm run efficiently within a 4GB VRAM envelope. diff --git a/Predict.py b/Predict.py new file mode 100644 index 0000000..554d887 --- /dev/null +++ b/Predict.py @@ -0,0 +1,67 @@ + + + +import torch +import numpy as np + +@torch.inference_mode() # More memory-efficient than no_grad() +def get_loss_per_sample(model, data_loader, device): + """ + Returns a list of individual losses for every sample in the loader. + Useful for MIA to see how 'certain' the model is about specific images. + """ + model.eval() + criterion = torch.nn.CrossEntropyLoss(reduction='none') # Crucial: returns loss per image + all_losses = [] + + for inputs, labels in data_loader: + inputs, labels = inputs.to(device), labels.to(device) + + outputs = model(inputs) + + # Calculate loss for each image in the batch individually + loss = criterion(outputs, labels) + + all_losses.extend(loss.cpu().numpy()) + + return all_losses + + +@torch.inference_mode() +def get_losses_by_class(model, data_loader, device): + """ + Returns a dictionary: { class_id: [list_of_losses_for_this_class] } + """ + model.eval() + criterion = torch.nn.CrossEntropyLoss(reduction='none') + + class_losses = {} + + for inputs, labels in data_loader: + inputs, labels = inputs.to(device), labels.to(device) + outputs = model(inputs) + + # Get individual losses + losses = criterion(outputs, labels).cpu().numpy() + labels_np = labels.cpu().numpy() + + for i, class_id in enumerate(labels_np): + if class_id not in class_losses: + class_losses[class_id] = [] + class_losses[class_id].append(losses[i]) + + return class_losses + + +# evaluate MIA +def eval_MIA(forgotten_losses, never_seen_losses): + avg_f_loss = np.mean(forgotten_losses) + avg_ns_loss = np.mean(never_seen_losses) + + print(f"Average Loss on Forgotten Identity: {avg_f_loss:.4f}") + print(f"Average Loss on Unknown Identities: {avg_ns_loss:.4f}") + + if avg_f_loss < avg_ns_loss * 0.8: + print("MIA Warning: Model still shows high certainty on forgotten data.") + else: + print("MIA Success: Model treats forgotten data like unknown data.") diff --git a/ReadME.md b/ReadME.md new file mode 100644 index 0000000..080db7d --- /dev/null +++ b/ReadME.md @@ -0,0 +1,85 @@ +# Python venv +Start a python environment here in this directory +```py +python -m venv . +``` + +Then we start the env using +```py +source ./bin/activate +``` + +We can then install whats needed with `pip`. for exampe +we can put all dependencies in some text file. say dependencies.txt +```py +# pip install +# already added dependencies.txt +pip install -r dependencies.txt + +``` + +Downloading the data from google drive was impossible. So Downloaded them manualy +and They need to be put in the a ./data directory +The download url was available in the error log. +`https://drive.google.com/uc?id=0B7EVK8r0v71pZjFTYXZWM3FlRnM` +this is the same location thats available in the official site + +``` +Root_dir/ +└── data/ + └── celeba/ + ├── img_align_celeba.zip + ├── list_attr_celeba.txt + ├── list_bbox_celeba.txt + ├── list_eval_partition.txt + └── list_landmarks_align_celeba.txt + +``` + +once this is manually done, We can run finetunning a selected model. For now, 8 models are implemented. +- ResNet-18 +- ResNet-50 +- DenseNet121 +- Inception +- GoogleNet +- ShuffleNet +- EfficientNet +- WideResNet + +## Fine tuning +### Preparation + +Lets say we want to finetune **Inception**. In `Tune.py` we have to adjust the variables accordingly like so: +```py +# Set the class size. e.g +CLASS_SIZE = 50 +# set the batch eg. +BATCH_SIZE = 16 +# set the Tuning epochs. e.g +EPOCH = 20 +# set the correct image size +# if ResNet or DenseNet, we set this to 224 +RESOLUTION = 299 +# set the model architecture +arch = Architecture.INCEPTION +``` +Other variable that we can change are those that are related to data size. Namely Training sample size and full sample size. +```py +# full sample size per class +SAMPLE_SIZE = 30 + +# Training sample size is then (full_sample - test_sample) +TRAINING_SMPLE = 28 + +# while at it, we can also set the learning rate +LR_RATE = 0.0001 +``` + + +### Rune the process +After we have set all necessary variables to our liking, we run the process by running Tune.py with python interpreter +```shell +# open terminal, cd to project root and run +python Tune.py +``` + diff --git a/SetUp.py b/SetUp.py new file mode 100644 index 0000000..939998c --- /dev/null +++ b/SetUp.py @@ -0,0 +1,14 @@ +## +import torch +from torchvision import datasets, transforms, models + +def get_device(): + + if torch.cuda.is_available(): + # clear cach to boost memory + # for new round + torch.cuda.empty_cache() + return torch.device("cuda") + else: + return torch.device("cpu") + diff --git a/Tune.py b/Tune.py new file mode 100644 index 0000000..3abb32d --- /dev/null +++ b/Tune.py @@ -0,0 +1,268 @@ +# Finetuning a selected model +# on a selected dataset +# using selected parameters + +from torch.utils.data import DataLoader +from sklearn.metrics import classification_report +import SetUp +#from Data import * +# from datasets.Casia import * +#from IdentitySubset import IdentitySubset +from sets.Data import * +from sets.IdentitySubset import IdentitySubset +# models +from architectures.Model import Model, Architecture + +from unlearning.LinearFiltration import LinearFiltration +from unlearning.CertifiedRemoval import CertifiedRemoval +from unlearning.WeightFiltration import WeightFiltration + +import Util +# WeightFiltration, CertifiedRemoval + +# numbre of classes +CLASS_SIZE = 20 +# batch +BATCH_SIZE = 16 + +# size of images per class trainset + testset +# 30 works best, more than that and we dont have enough data +SAMPLE_SIZE = 30 + +# this is then (full_sample - test_sample) +TRAINING_SMPLE = 27 + +# learning rate +LR_RATE = 0.0001 +EPOCHS = 10 + +# depends on model architecture +# ResNet, DenseNet = 224 +# Inception = 299 +RESOLUTION = 224 + +FINETUNE = False # whether to fintune or just load finetuned model from dir +# model architecture options are +# - RESNET18 +# - RESNET50 +# - DENSENET121 +# - INCEPTION +# - GOOGLENET +# - EFFICIENTNET +# - SHUFFLENET +arch = Architecture.RESNET50 + +# DATA PREPARATION +# load data set and prepare +dataset_name = Set_Name.CELEBA +set = Set_Name.CELEBA + +dataset = get_set(set_name=dataset_name) +print(f"> {dataset.__class__.__name__} dataset loaded") +# select identities for experiment +#selected_identities = select_ids( +# dataset = dataset, +# sample_size = SAMPLE_SIZE, +# class_size = CLASS_SIZE +# ) + +# this selects the top 50 based on sample size +# that way repeated calls return the same classes +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_SMPLE} train and {SAMPLE_SIZE-TRAINING_SMPLE} test samples') + +# split class images to train/test indices +train_indices, test_indices = get_indices( + dataset = dataset, + identities = selected_identities, + split_at = TRAINING_SMPLE, + size= SAMPLE_SIZE +) + +# helps map class id to index +id_map = {old_id: new_id for new_id, old_id in enumerate(selected_identities)} + +# we remap identities because crossEntropyLoss requires in indices 0 -> (n-1) +# where n = class size. +tr_transform = train_transform(RESOLUTION) +train_data = IdentitySubset( + dataset=dataset, + indices=train_indices, + id_mapping=id_map, + transform=tr_transform) + +train_loader = DataLoader( + train_data, + batch_size = BATCH_SIZE, + shuffle = True) + +print(f"> Total training images: {len(train_data)}") + +print(f'> Constants : Classes = {CLASS_SIZE}, Batch = {BATCH_SIZE}, epochs = {EPOCHS}') + +# MODEL PREPARATION +# cuda if exists (it does here) +device = SetUp.get_device() + + +for i in range(0,1):#CLASS_SIZE): + FORGET_CLASS_IDX = i + # Create model using Factory + + model = None + + 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 + + # EVALUATE + te_transform = test_transform(RESOLUTION) + # Testing + test_data = IdentitySubset( + dataset = dataset, + indices=test_indices, + id_mapping=id_map, + transform=te_transform) + + test_loader = DataLoader( + test_data, + batch_size=BATCH_SIZE, + shuffle=False) + + print(f"Total test images for these {CLASS_SIZE} classes: {len(test_data)}") + + # Evaluate + current_mode = "Finetuned" + if FINETUNE: + + #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) + #filtration.apply(reloaded.model) + + #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, + 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 + forget_train_loader, retain_train_loader = get_unlearning_loaders( + dataset=train_data, + forget_class_idx=FORGET_CLASS_IDX, + batch_size=BATCH_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 + ) + + + #strategies = [linear_filtration, weight_filtration, certified_removal] + strategies = [certified_removal] + for strategy in strategies: + # test again + reloaded = Model.create( + arch=arch, + device = device, + size = CLASS_SIZE + ) + reloaded.load(arch = arch) + print("fine tunned model loaded") + # reloaded.evaluate( + # 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) + # Performance Analysis + strategy_in_use = strategy.__class__.__name__ + + # evaluation on retain Test_set + current_mode = "retain" + print("\n--- Performance on Retained Classes") + accuracy, report_dict = reloaded.evaluate(loader=retain_test_loader, mode=current_mode) + + Util._log_to_csv( + arch=reloaded.__class__.__name__, + mode = current_mode, + accuracy=accuracy, + report_dict=report_dict, + strategy=strategy_in_use + ) + + + # evaluation on forget Test_set + print("\n--- Performance on Forgotten Class") + current_mode = "forget" + accuracy, report_dict = reloaded.evaluate(loader=forget_test_loader,mode=current_mode) + Util._log_to_csv( + arch=reloaded.__class__.__name__, + mode = current_mode, + accuracy=accuracy, + report_dict=report_dict, + strategy=strategy_in_use + ) + + # evaluation on forget Train_set + # we expect this to be equal or highr than accuracy on Forget Test_set + current_mode = "forget_train" + print("\n--- Performance on Forgotten Class (Train Set - Verifying Unlearning)") + accuracy, report_dict = reloaded.evaluate(loader=forget_train_loader, mode=current_mode) + Util._log_to_csv( + arch= reloaded.__class__.__name__, + mode = current_mode, + accuracy=accuracy, + report_dict=report_dict, + strategy=strategy_in_use + ) \ No newline at end of file diff --git a/Tune_new.py b/Tune_new.py new file mode 100644 index 0000000..0998593 --- /dev/null +++ b/Tune_new.py @@ -0,0 +1,285 @@ +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.CertifiedUnlearning import CertifiedUnlearning +from unlearning.LinearFiltration import LinearFiltration +from unlearning.WeightFiltration import WeightFiltration + + +# Global Hyperparameters +CLASS_SIZE = 20 +BATCH_SIZE = 16 +SAMPLE_SIZE = 30 +TRAINING_SAMPLE = 27 + +# depends on model architecture +# ResNet, DenseNet = 224 +# Inception = 299 +RESOLUTION = 224 + +# specify the model architecture, +# Options here are the following +''' + RESNET18 # candidate + RESNET50 + RESNET34 + INCEPTION # candidate / or googleNet + DENSENET121 # candidate + GOOGLENET # candidate / or Inception + EFFICIENTNET # candidate + SHUFFLENET + WIDE_RESNET +''' +ARCH = Architecture.RESNET34 + + +# 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.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 {dataset_name.name} 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=14): + + + """ + 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) + + + train_loader = DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True) + + if not run_training: + return + + # Finetuning + 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=ARCH.name,#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, strategy, evaluate = False): + """ + 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"] + + + # Segment specific unlearning loaders using class index boundaries + retain_train_loader , forget_train_loader= get_unlearning_loaders( + dataset=train_data, forget_class_idx=forget_class_idx, batch_size=BATCH_SIZE + ) + retain_test_loader, forget_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) + + if evaluate: + reloaded.evaluate( + loader=retain_test_loader, mode="finetuned" + ) + + print("fine tunned model loaded into evaluation sandbox") + + # Execute strategic parameter unlearning step + 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"}, + {"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=ARCH.name,#reloaded.__class__.__name__, + mode=domain["mode"], + accuracy=accuracy, + report_dict=report_dict, + strategy=strategy_in_use + ) + + +# entry +if __name__ == "__main__": + + outer_loop = 1 + inner_loop = CLASS_SIZE + + for k in range(outer_loop): + + try: + # Data Infrastructure and Architecture + 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) + # scale 16400.0 for ResNet + scale = 22000 + # batch 8 for resNet, + unlearning_batches = 32 + # regularis + # strategies + certified_unlearning = CertifiedUnlearning( + target_class_index=0, #arch ResNet18 GoogLeNet Inception + l2_reg=0.000002 , # 0.000002 0.00001 0.0 + gamma=0.01, # 0.1 0.1 0.01 + scale= scale, # 16400.0 35000.0 + s1=2, # 2 + s2=150, # 300 + std=0.00001, # 0.00001 + unlearn_bs=unlearning_batches # 8 32 8 + ) + + # works perfectly + linear_filtration = LinearFiltration( + + target_class_index=0 + ) + + weight_filtration = WeightFiltration( + target_class_index=0, #arch ResNet18 GoogLeNet/Inception + epochs=6, # + lr=250.0, # ResNet18 = 150 # 150 100 + gamma=0.001, # 0.001 + lambda_1=30, # 25 100 + arch=ARCH + ) + + strategies = [ + certified_unlearning, + #weight_filtration, + #linear_filtration + ] + # Unlearning Iteration + for i in range(0, inner_loop): + + 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, + # if we are finetuning, no need to evaluate base model. + # or may be never when not either! + evaluate = not finetuning + ) + + except KeyboardInterrupt: + print("program interrupted. Exit!") + break diff --git a/Util.py b/Util.py new file mode 100644 index 0000000..2910054 --- /dev/null +++ b/Util.py @@ -0,0 +1,51 @@ + +from pathlib import Path +import time +import os + +def _log_to_csv(arch, mode, accuracy, report_dict, strategy): + """Handles directory structures, file setups, and distinct CSV column formatting.""" + #arch_name = model.__class__.__name__.lower() + save_dir = Path(f"reports/{strategy}/{arch}") + save_dir.mkdir(parents=True, exist_ok=True) + csv_path = save_dir / f"{mode}.csv" + + file_exists = csv_path.exists() + + headers = [ + "accuracy", + "macro_precision", "macro_recall", "macro_f1", + "weighted_precision", "weighted_recall", "weighted_f1" + ] + row = [ + f"{accuracy / 100.0:.4f}", + f"{report_dict['macro avg']['precision']:.4f}", + f"{report_dict['macro avg']['recall']:.4f}", + f"{report_dict['macro avg']['f1-score']:.4f}", + f"{report_dict['weighted avg']['precision']:.4f}", + f"{report_dict['weighted avg']['recall']:.4f}", + f"{report_dict['weighted avg']['f1-score']:.4f}" + ] + + with open(csv_path, "a") as f: + if not file_exists: + f.write(",".join(headers) + "\n") + f.write(",".join(row) + "\n") + + print(f">> Direct CSV metrics appended to {csv_path}") + + +def _initialize_log_file(log_file): + """Creates a unique log file for this strategy with a header if it doesn't exist.""" + log_file.parent.mkdir(parents=True, exist_ok=True) + if not os.path.exists(log_file): + with open(log_file, "w") as f: + f.write("execution_time_sec\n") + +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") + + + diff --git a/architectures/DenseNet121.py b/architectures/DenseNet121.py new file mode 100644 index 0000000..68a6735 --- /dev/null +++ b/architectures/DenseNet121.py @@ -0,0 +1,15 @@ + +import torch.nn as nn +from torchvision import models +from architectures.Model import Model + +class DenseNet121(Model): + def get(self): + + # load pretrained + m = models.densenet121(weights=models.DenseNet121_Weights.DEFAULT) + # will modify only the final layers + num_ftrs = m.classifier.in_features + m.classifier = nn.Linear(num_ftrs, self.size) + + return m \ No newline at end of file diff --git a/architectures/EfficentNet.py b/architectures/EfficentNet.py new file mode 100644 index 0000000..9fbfc97 --- /dev/null +++ b/architectures/EfficentNet.py @@ -0,0 +1,19 @@ + +import torch.nn as nn +from torchvision import models + +# Base model +from architectures.Model import Model + +class EfficientNet(Model): + + def get(self): + + m = models.efficientnet_b1(weights=models.EfficientNet_B1_Weights.DEFAULT) + + # Unfreeze the last block for a lighter touch + for param in m.features[-1].parameters(): param.requires_grad = True + + # Standard classifier fix + m.classifier[1] = nn.Linear(m.classifier[1].in_features, self.size) + return m \ No newline at end of file diff --git a/architectures/GoogleNet.py b/architectures/GoogleNet.py new file mode 100644 index 0000000..8dd9c9c --- /dev/null +++ b/architectures/GoogleNet.py @@ -0,0 +1,31 @@ + +import torch.nn as nn +from torchvision import models + +# Base model +from architectures.Model import Model + +class GoogleNet(Model): + + def get(self): + + m = models.googlenet(weights=models.GoogLeNet_Weights.DEFAULT) + + # 1. Handle the two Auxiliary Classifiers + # GoogLeNet has aux1 and aux2 to help training converge + #if m.aux_logits: + #m.aux1.fc = nn.Linear(m.aux1.fc.in_features, self.size) + #m.aux2.fc = nn.Linear(m.aux2.fc.in_features, self.size) + + # 2. Handle the Main Classifier + m.fc = nn.Linear(m.fc.in_features, self.size) + + #for param in m.parameters(): + # param.requires_grad = False + + # Unfreezing the final stages for identity recognition + #for name, param in m.named_parameters(): + # if "inception5" in name or "fc" in name: + # param.requires_grad = True + + return m \ No newline at end of file diff --git a/architectures/Inception.py b/architectures/Inception.py new file mode 100644 index 0000000..b55997f --- /dev/null +++ b/architectures/Inception.py @@ -0,0 +1,47 @@ + +import torch +import torch.nn as nn +import torch.optim as optim +from torchvision import models +import time + +# Base model +from architectures.Model import Model + +class Inception(Model): + def get(self): + m = models.inception_v3(weights=models.Inception_V3_Weights.DEFAULT) + #for param in model.parameters(): + # param.requires_grad = False + m.AuxLogits.fc = nn.Linear(m.AuxLogits.fc.in_features, self.size) + m.fc = nn.Linear(m.fc.in_features, self.size) + return m + + def train(self, epochs, loader, rate): + # Override because Inception returns a tuple (main, aux) + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(filter(lambda p: p.requires_grad, self.model.parameters()), lr=rate) + + print(f"Starting training on {self.device}...") + start_time = time.time() + + self.model.train() + + for epoch in range(epochs): + total_loss = 0.0 + for inputs, labels in loader: + + inputs, labels = inputs.to(self.device), labels.to(self.device) + optimizer.zero_grad() + + outputs, aux_outputs = self.model(inputs) + loss = criterion(outputs, labels) + 0.3 * criterion(aux_outputs, labels) + + loss.backward() + optimizer.step() + total_loss += loss.item() + + print(f"Epoch {epoch+1}/{epochs} | Loss: {total_loss/len(loader):.4f}") + + if self.device.type == 'cuda': torch.cuda.synchronize() + print(f"Training completed in: {time.time() - start_time:.2f}s") diff --git a/architectures/Model.py b/architectures/Model.py new file mode 100644 index 0000000..855c4c6 --- /dev/null +++ b/architectures/Model.py @@ -0,0 +1,225 @@ + +from abc import ABC, abstractmethod +import torch +import torch.nn as nn +import torch.optim as optim +import time +import numpy as np +from sklearn.metrics import classification_report +from pathlib import Path +#from unlearning.Strategy import Strategy +import copy +from torch.optim.lr_scheduler import CosineAnnealingLR +import Util + +class Model(ABC): + # need to add a weight decay here + def __init__(self, device, size): + self.device = device + self.size = size + self.model = self.get().to(self.device) + + @abstractmethod + def get(self): + pass + + ''' + Have to have a new param here as mode, for example it would be base, or retrain + param mode = "base" or "retrain" + that way I can save time it takes to train and retrain. + file would be solved with Util functions + log_file = Path(f"reports/{mode}/{self.__class__.__name__}/time_metrics.txt") + Util._initialize_log_file(log_file): + strt = time.perf_counter() + end = time.perf_counter() + and then save logs + execution_time = end -strt + Util.log_metric(log_file, execution_time: float): + + ''' + def train(self, epochs, loader, rate, mode = "retrain"): + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(filter(lambda p: p.requires_grad, self.model.parameters()), lr=rate, weight_decay=0.1) + + scheduler = CosineAnnealingLR(optimizer, T_max=epochs, eta_min=1e-6) + # to save reports + file_path = Path(f"{mode}/{self.__class__.__name__.lower()}/time_metrics.txt") + Util._initialize_log_file(file_path) + + #save_dir.mkdir(parents=True, exist_ok=True) + + print(f"Starting training on {self.device}...") + start_time = time.time() + self.model.train() + + for epoch in range(epochs): + total_loss = 0.0 + for inputs, labels in loader: + inputs, labels = inputs.to(self.device), labels.to(self.device) + optimizer.zero_grad() + outputs = self.model(inputs) + loss = criterion(outputs, labels) + loss.backward() + optimizer.step() + total_loss += loss.item() + scheduler.step() + + print(f"Epoch {epoch+1}/{epochs} | Loss: {total_loss / len(loader):.4f}") + end_time = time.time() + Util.log_metric(log_file=file_path, execution_time=(end_time - start_time)) + if self.device.type == 'cuda': torch.cuda.synchronize() + print(f"Training completed in: {time.time() - start_time:.2f}s") + + + def save(self, filename=None): + + save_dir = Path("trained_models") + save_dir.mkdir(parents=True, exist_ok=True) + + # Filename (Default to class name if not provided) + if filename is None: + filename = f"{self.__class__.__name__.lower()}.pth" + + if not filename.endswith('.pth'): + filename += '.pth' + + save_path = save_dir / filename + torch.save(self.model.state_dict(), save_path) + print(f'Model saved to {save_path}') + + + + def load(self, arch): + + file_path = Path("trained_models") / f'{arch.name.lower()}.pth' + # does file exist + if not file_path.exists(): + raise FileNotFoundError(f'No checkpoint found at: {file_path}') + + # Load the weights + state_dict = torch.load(file_path, map_location=self.device, weights_only=True) + + self.model.load_state_dict(state_dict) + self.model.to(self.device) + print(f'Model loaded from {file_path}') + + + def unlearn(self, strategy: 'Strategy', forget_loader, retain_loader): + """ Executes a targeted unlearning strategy and profiles efficiency """ + print(f"Executing: {strategy.__class__.__name__}...") + + start_time = time.time() + + # Delegate the actual algorithmic weight/logit manipulation to the strategy + strategy.apply(self.model, forget_loader, retain_loader) + + elapsed_time = time.time() - start_time + print(f"{strategy.__class__.__name__} completed in {elapsed_time:.4f} seconds.") + + return elapsed_time + + def evaluate(self, loader, mode="eval"): + """ + 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}]...") + + with torch.no_grad(): + for inputs, labels in loader: + inputs, labels = inputs.to(self.device), labels.to(self.device) + outputs = self.model(inputs) + _, predicted = torch.max(outputs, 1) + all_preds.extend(predicted.cpu().numpy()) + all_labels.extend(labels.cpu().numpy()) + + # Extract only the active classes evaluated in this loader slice + classes = sorted(list(set(all_labels))) + accuracy = 100 * (np.array(all_preds) == np.array(all_labels)).sum() / len(all_labels) + + print(f"Test Accuracy: {accuracy:.2f}%") + + # 1. Print standard text report to terminal + print(classification_report(all_labels, all_preds, labels=classes, zero_division=0)) + + # 2. Extract structured dictionary metrics + report_dict = classification_report( + all_labels, + all_preds, + labels=classes, + output_dict=True, + zero_division=0 + ) + + # 3. Delegate file tracking to isolated helper method + #self._log_to_csv(mode, accuracy,report_dict) + return accuracy, report_dict + + + + # factory + @staticmethod + def create(arch, device, size): + print(f'>> MODEL ARCHITECTURE >> {arch.name}.') + + match arch: + + # ResNet18 + case Architecture.RESNET18: + 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 + return ResNet50(device, size) + + # INCEPTION + case Architecture.INCEPTION: + from architectures.Inception import Inception + return Inception(device, size) + + # DENSENET121 + case Architecture.DENSENET121: + from architectures.DenseNet121 import DenseNet121 + return DenseNet121(device, size) + # googleNet + case Architecture.GOOGLENET: + from architectures.GoogleNet import GoogleNet + return GoogleNet(device, size) + # EfficientNet + case Architecture.EFFICIENTNET: + from architectures.EfficentNet import EfficientNet + return EfficientNet(device, size) + #ShuffleNet + case Architecture.SHUFFLENET: + from architectures.ShuffleNet import ShuffleNet + return ShuffleNet(device, size) + # wide ResNet + case Architecture.WIDE_RESNET: + from architectures.WideResNet import WideResNet + return WideResNet(device, size) + case _: + raise ValueError(f"Unknown model: {arch}") + +# model architectures +from enum import Enum, auto + +class Architecture(Enum): + RESNET18 = auto() + RESNET50 = auto() + RESNET34 = auto() + INCEPTION = auto() + DENSENET121 = auto() + GOOGLENET = auto() + EFFICIENTNET = auto() + SHUFFLENET = auto() + WIDE_RESNET = auto() \ No newline at end of file diff --git a/architectures/ResNet18.py b/architectures/ResNet18.py new file mode 100644 index 0000000..84cb19e --- /dev/null +++ b/architectures/ResNet18.py @@ -0,0 +1,22 @@ + +import torch.nn as nn +from torchvision import models + +# Base model +from architectures.Model import Model + +class ResNet18(Model): + + def get(self): + m = models.resnet18(weights=models.ResNet18_Weights.DEFAULT) + + # freez all layers + #for param in m.parameters(): + # param.requires_grad = False + + # unfreez the last two + #for param in m.layer3.parameters(): param.requires_grad = True + #for param in m.layer4.parameters(): param.requires_grad = True + + m.fc = nn.Linear(m.fc.in_features, self.size) + return m \ No newline at end of file diff --git a/architectures/ResNet34.py b/architectures/ResNet34.py new file mode 100644 index 0000000..19eea6c --- /dev/null +++ b/architectures/ResNet34.py @@ -0,0 +1,22 @@ + +import torch.nn as nn +from torchvision import models + +# Base model +from architectures.Model import Model + +class ResNet34(Model): + + def get(self): + m = models.resnet34(weights=models.ResNet34_Weights.DEFAULT) + + # freez all layers + #for param in m.parameters(): + # param.requires_grad = False + + # unfreez the last two + #for param in m.layer3.parameters(): param.requires_grad = True + #for param in m.layer4.parameters(): param.requires_grad = True + + m.fc = nn.Linear(m.fc.in_features, self.size) + return m \ No newline at end of file diff --git a/architectures/ResNet50.py b/architectures/ResNet50.py new file mode 100644 index 0000000..860271b --- /dev/null +++ b/architectures/ResNet50.py @@ -0,0 +1,36 @@ + +import torch.nn as nn +from torchvision import models + +# Base model +from architectures.Model import Model + +class ResNet50(Model): + # NOTE: + # This model had it's best performance with the following configs + # numbre of classes + # CLASS_SIZE = 20 + # BATCH_SIZE = 16 + # SAMPLE_SIZE = 30 + # TRAINING_SMPLE = 28 + # LR_RATE = 0.0001 + # EPOCHS = 15 + # RESOLUTION = 224 + # NOTE: But it may be a one time thing. + # because testing again didn't repeat + + def get(self): + m = models.resnet50(weights=models.ResNet50_Weights.DEFAULT) + + # freez all layers + #for param in m.parameters(): + #param.requires_grad = False + + # unfreez the last two + # NOTE: Freezing everything and unfrizing the last 3 yeilded the best performance + #for param in m.layer2.parameters(): param.requires_grad = True + #for param in m.layer3.parameters(): param.requires_grad = True + #for param in m.layer4.parameters(): param.requires_grad = True + + m.fc = nn.Linear(m.fc.in_features, self.size) + return m \ No newline at end of file diff --git a/architectures/ShuffleNet.py b/architectures/ShuffleNet.py new file mode 100644 index 0000000..fc2a52d --- /dev/null +++ b/architectures/ShuffleNet.py @@ -0,0 +1,17 @@ + + + +import torch.nn as nn +from torchvision import models + +# Base model +from architectures.Model import Model + +class ShuffleNet(Model): + + def get(self): + m = models.shufflenet_v2_x1_0(weights=models.ShuffleNet_V2_X1_0_Weights.DEFAULT) + + num_ftrs = m.fc.in_features + m.fc = nn.Linear(num_ftrs, self.size) + return m \ No newline at end of file diff --git a/architectures/WFNet.py b/architectures/WFNet.py new file mode 100644 index 0000000..b4a20df --- /dev/null +++ b/architectures/WFNet.py @@ -0,0 +1,230 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torch.utils.data import DataLoader +import numpy as np +from sklearn.metrics import classification_report +from architectures.Model import Model + +'''class WF_Module(nn.Module): + """ + Pure PyTorch Neural Network module graph. + Keeps parameter registration and autograd tracking separate from + the framework's high-level Model abstractions to prevent recursion collisions. + """ + def __init__(self, original_model: nn.Module, num_classes: int): + super().__init__() + + self.original_model = original_model + + # Target layer for weight filtering (layer4 block 1 conv2 or conv3 depending on arch) + last_layer = original_model.layer4[1] + + # Some versions are limited to 2 convolutional layers + if hasattr(last_layer, "conv3"): + self.target_conv = last_layer.conv3 + else: + self.target_conv = last_layer.conv2 + + # Completely freeze the original ResNet parameters + for param in self.parameters(): + param.requires_grad = False + + # Initialize the alpha parameter matrix (Rows = Classes, Cols = Channels) + out_channels = self.target_conv.weight.shape[0] + self.alpha = nn.Parameter(torch.full((num_classes, out_channels), 3.0))''' + +''' + Poppi et_al's Single-shot multiclass unlearning. + This calculation happens only once to generate the mask. once the mask is generated, + Unlearning and remembering becomes a matter of switching gates on and off.''' +''' + def forward(self, x: torch.Tensor, target_class_indices: torch.Tensor) -> torch.Tensor: + # we linearly loop through layers 1 to 4[block 1] (for ResNet) + # for i in M_{|L|} do l <- l[i] + x = self.original_model.maxpool(self.original_model.relu(self.original_model.bn1(self.original_model.conv1(x)))) + x = self.original_model.layer1(x) + x = self.original_model.layer2(x) + x = self.original_model.layer3(x) + x = self.original_model.layer4[0](x) + + # The second block execute its internal transformations natively + # This handles conv1->conv2 (ResNet18) or conv1->conv2->conv3 (ResNet50) automatically! + # Xi+1 <- l(Xi, ˆwl) + x = self.original_model.layer4[1](x) + + # Apply mask dynamically to the completed block feature map + # wl <- αl[Yunl] ⊙ ˆwl + batch_alpha = self.alpha[target_class_indices] + mask = torch.sigmoid(batch_alpha).view(x.size(0), -1, 1, 1) + x = x * mask + + # Remaining standard head steps + x = self.original_model.avgpool(x) + x = torch.flatten(x, 1) + # so here we are returning the output logits + # the result of classification is then + # argmax(x) + return self.original_model.fc(x) +''' + +class WF_Module(nn.Module): + def __init__(self, original_model: nn.Module, num_classes: int, arch_enum): + super().__init__() + # If your model classes contain the raw inner torch model under an attribute, + # extract it. Otherwise, use it directly. + self.original_model = getattr(original_model, "model", original_model) + + # Freeze the original model parameters completely + for param in self.original_model.parameters(): + param.requires_grad = False + + # Target layer discovery using your clean Enum contract + self.target_layer = self._deduce_target_layer(self.original_model, arch_enum) + + # Derive channel dimensions dynamically from the deduced layer + out_channels = self._extract_channels(self.target_layer, self.original_model) + + # Initialize alpha parameter matrix (Rows = Classes, Cols = Channels) + self.alpha = nn.Parameter(torch.full((num_classes, out_channels), 3.0)) + self._current_target_indices = None + + def _deduce_target_layer(self, model: nn.Module, arch_enum) -> nn.Module: + """ + Scans the architecture topology to target the final deep feature extraction block + right before global pooling/classification using strict Enum configurations. + """ + match arch_enum: + # --- RESNET FAMILY --- + case arch_enum.RESNET18 | arch_enum.RESNET34 | arch_enum.RESNET50 | arch_enum.WIDE_RESNET: + return model.layer4[-1] + + # --- GOOGLENET --- + case arch_enum.GOOGLENET: + return model.inception5b + + # --- INCEPTION V3 --- + case arch_enum.INCEPTION: + return model.Mixed_7c + + # --- DENSENET 121 --- + case arch_enum.DENSENET121: + return model.features.norm5 + + # --- EFFICIENTNET --- + case arch_enum.EFFICIENTNET: + return model.features[-1] + + # --- SHUFFLENET --- + case arch_enum.SHUFFLENET: + return model.conv5 + + case _: + # Robust Fallback Strategy + target = None + for module in model.modules(): + if isinstance(module, nn.Conv2d): + target = module + if target is not None: + return target + raise RuntimeError(f"Could not locate filtration anchor for Enum target: {arch_enum}") + + def _extract_channels(self, target_layer: nn.Module, model: nn.Module) -> int: + """Helper to determine channel depth across varied layers types.""" + if hasattr(target_layer, "out_channels"): + return target_layer.out_channels + if hasattr(target_layer, "num_features"): + return target_layer.num_features + if hasattr(target_layer, "weight"): + return target_layer.weight.shape[0] + + # Classifier fallback mapping + if hasattr(model, "fc"): + return model.fc.in_features + if hasattr(model, "classifier"): + if isinstance(model.classifier, nn.Linear): + return model.classifier.in_features + if isinstance(model.classifier, nn.Sequential): + return model.classifier[0].in_features + return 512 + + def _filtration_hook(self, module: nn.Module, hook_input: tuple, hook_output: torch.Tensor) -> torch.Tensor: + if self._current_target_indices is None: + return hook_output + + batch_alpha = self.alpha[self._current_target_indices] + + if len(hook_output.shape) == 4: + mask = torch.sigmoid(batch_alpha).view(hook_output.size(0), -1, 1, 1) + else: + mask = torch.sigmoid(batch_alpha).view(hook_output.size(0), -1) + + return hook_output * mask + + def forward(self, x: torch.Tensor, target_class_indices: torch.Tensor) -> torch.Tensor: + self._current_target_indices = target_class_indices + hook_handle = self.target_layer.register_forward_hook(self._filtration_hook) + try: + logits = self.original_model(x) + finally: + hook_handle.remove() + self._current_target_indices = None + return logits + + +class WF_Net_Model(Model): + def __init__(self, device, size, original_model: nn.Module, target_class_index: int, arch): + self.device = device + self.size = size + self.wf_module = WF_Module( + arch_enum=arch, + original_model = original_model, + num_classes = size + ).to(self.device) + + # this index indicates which row of the mask should be active (gate closed). + self.target_class_index = target_class_index + self.model = self.wf_module + + def get(self): + return self.wf_module + + ''' + We override the evaluate method from the base class, + because how we evaluate is different here from that of a normal torch nn.Module object + + ''' + def evaluate(self, loader, mode="eval"): + + self.wf_module.eval() + all_preds, all_labels = [], [] + print(f"\nEvaluating Domain: [{mode}]...") + + with torch.no_grad(): + for inputs, labels in loader: + inputs, labels = inputs.to(self.device), labels.to(self.device) + + # we apply the filter + gate_signals = torch.full((inputs.size(0),), self.target_class_index, dtype=torch.long, device=self.device) + + # pass prediction through the filter + outputs = self.wf_module(inputs, target_class_indices=gate_signals) + + # return argmax(x) + _, predicted = torch.max(outputs, 1) + all_preds.extend(predicted.cpu().numpy()) + all_labels.extend(labels.cpu().numpy()) + + classes = sorted(list(set(all_labels))) + accuracy = 100 * (np.array(all_preds) == np.array(all_labels)).sum() / len(all_labels) + + print(f"Test Accuracy: {accuracy:.2f}%") + print(classification_report(all_labels, all_preds, labels=classes, zero_division=0)) + report = classification_report(all_labels, all_preds, labels=classes, output_dict=True, zero_division=0) + + return accuracy, report + + def eval(self): + """Safely intercept any fallback base class calls targeting .eval()""" + self.wf_module.eval() diff --git a/architectures/WideResNet.py b/architectures/WideResNet.py new file mode 100644 index 0000000..b645096 --- /dev/null +++ b/architectures/WideResNet.py @@ -0,0 +1,15 @@ + + +import torch.nn as nn +from torchvision import models + +# Base model +from architectures.Model import Model + +class WideResNet(Model): + + def get(self): + # wide_resnet50_2 is a common high-performance choice + m = models.wide_resnet50_2(weights=models.Wide_ResNet50_2_Weights.DEFAULT) + m.fc = nn.Linear(m.fc.in_features, self.size) + return m \ No newline at end of file diff --git a/dependencies.txt b/dependencies.txt new file mode 100644 index 0000000..e2c6522 --- /dev/null +++ b/dependencies.txt @@ -0,0 +1,6 @@ +torch +torchvision +gdown +numpy +scikit-learn +kagglehub \ No newline at end of file diff --git a/js_evaluator/JS_Evaluator.py b/js_evaluator/JS_Evaluator.py new file mode 100644 index 0000000..3961e9b --- /dev/null +++ b/js_evaluator/JS_Evaluator.py @@ -0,0 +1,111 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.utils.data import DataLoader +import torchvision.models as models + +class ZeroRetrainForgettingEvaluator: + def __init__(self, unlearned_model: nn.Module, num_classes: int): + """ + Initializes the ZRF Evaluator. + + Args: + unlearned_model (nn.Module): Your fine-tuned & unlearned ResNet-50. + num_classes (int): Number of classes used in your CelebA task. + """ + # select device + if torch.cuda.is_available(): + self.device = torch.device("cuda") + elif hasattr(torch, "xpu") and torch.xpu.is_available(): + self.device = torch.device("xpu") # For Intel GPUs using IPEX + else: + self.device = torch.device("cpu") + + print(f"[INFO] Using device: {self.device}") + + # prepare the unlearned model + self.unlearned_model = unlearned_model.to(self.device) + self.unlearned_model.eval() + + # Instantiate a structurally matching, completely random model + print(f"[INFO] Initializing random baseline ResNet-50 with {num_classes} classes...") + self.random_model = self.get_random_model(num_classes) + self.random_model = self.random_model.to(self.device) + self.random_model.eval() + + # gets randomly initialised model + # for comparison with unlearned model + def get_random_model(num_classes): + print(f"[INFO] Initializing random baseline ResNet-50 with {num_classes} classes...") + model = models.resnet50(weights=None) + model.fc = nn.Linear(model.fc.in_features, num_classes) + return model + + + # compute divergence + def _compute_js_divergence(self, p: torch.Tensor, q: torch.Tensor) -> float: + """ + Computes the Jensen-Shannon (JS) Divergence between two probability distributions. + + Args: + p, q (Tensor): Tensors of shape (batch_size, num_classes) containing probabilities. + """ + # Avoid log(0) issues by adding a tiny epsilon + eps = 1e-12 + p = torch.clamp(p, eps, 1.0) + q = torch.clamp(q, eps, 1.0) + + # Calculate the midpoint distribution + m = 0.5 * (p + q) + + # Compute KL Divergence natively: KL(P || M) and KL(Q || M) + kl_pm = torch.sum(p * (torch.log(p) - torch.log(m)), dim=1) + kl_qm = torch.sum(q * (torch.log(q) - torch.log(m)), dim=1) + + # JS Divergence is the average of both KL divergences + js_div = 0.5 * (kl_pm + kl_qm) + + # Return the mean divergence across the entire batch + return js_div.mean().item() + + def evaluate_forget_class(self, dataset, batch_size: int = 32) -> float: + """ + Evaluates the unlearned model against the random model using images + from the forgotten class/identity. + + Args: + dataset (Dataset): A PyTorch Dataset containing images of the forget set. + batch_size (int): Batch size for evaluation. + + Returns: + float: The ZRF score (JS Divergence). A lower divergence means + the unlearned model is behaving exactly like a random model. + """ + dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False) + + total_js_div = 0.0 + total_samples = 0 + + # No gradients needed for evaluation + with torch.no_grad(): + for images, _ in dataloader: + images = images.to(self.device) + batch_len = images.size(0) + + # Get raw outputs (logits) + unlearned_logits = self.unlearned_model(images) + random_logits = self.random_model(images) + + # Convert logits to probability distributions via Softmax + unlearned_probs = F.softmax(unlearned_logits, dim=1) + random_probs = F.softmax(random_logits, dim=1) + + # Calculate JS divergence for this batch + batch_js = self._compute_js_divergence(unlearned_probs, random_probs) + + # Weighted average based on batch size (handles final smaller batches perfectly) + total_js_div += batch_js * batch_len + total_samples += batch_len + + final_zrf_score = total_js_div / total_samples + return final_zrf_score \ No newline at end of file diff --git a/sets/Casia.py b/sets/Casia.py new file mode 100644 index 0000000..0c8e87c --- /dev/null +++ b/sets/Casia.py @@ -0,0 +1,167 @@ + + +from torchvision import datasets, transforms +from torch.utils.data import Dataset, DataLoader, Subset +import torch +import numpy as np +import os + +# train set transform +def train_transform(res): + return transforms.Compose([ + transforms.Resize((res, res)), + transforms.RandomHorizontalFlip(p=0.5), + transforms.ColorJitter( + brightness=0.2, + contrast=0.2, + saturation=0.1 + ), + transforms.ToTensor(), + transforms.Normalize( + mean=[0.485, 0.456, 0.406], + std=[0.229, 0.224, 0.225] + ) + ]) + +# test set transform +def test_transform(res): + return transforms.Compose([ + transforms.Resize((res, res)), + transforms.ToTensor(), + transforms.Normalize( + mean=[0.485, 0.456, 0.406], + std=[0.229, 0.224, 0.225] + ) + ]) + +# Load data using ImageFolder for CASIA-WebFace +''' +def get_set(): + # This will check local cache first, then download if missing + print("Checking for CASIA-WebFace dataset...") + path = kagglehub.dataset_download("debarghamitraroy/casia-webface") + + # Kagglehub often downloads a nested structure (e.g., path/casia-webface/casia-webface) + # We need the folder that directly contains the identity subfolders + # We'll check if there's a 'casia-webface' subfolder inside the downloaded path + sub_path = os.path.join(path, "casia-webface") + final_path = sub_path if os.path.exists(sub_path) else path + + print(f"Loading dataset from: {final_path}") + + return datasets.ImageFolder( + root=final_path, + transform=None + )''' +# Load data using ImageFolder for your UNPACKED images +def get_set(): + # This must point to the folder created by Extractor.py + # NOT the kagglehub cache path + final_path = os.path.abspath("./data/casia-set") + + if not os.path.exists(final_path): + raise FileNotFoundError( + f"Unpacked dataset not found at {final_path}. " + "Please run Extractor.py first!" + ) + + print(f"Loading unpacked CASIA dataset from: {final_path}") + + return datasets.ImageFolder( + root=final_path, + transform=None + ) + +def get_ids_and_counts(dataset): + # ImageFolder stores labels in .targets + targets = torch.tensor(dataset.targets) + return torch.unique( + input = targets, + return_counts=True + ) + +def select_ids(dataset, sample_size, class_size): + ids, counts = get_ids_and_counts(dataset=dataset) + eligible_mask = counts >= sample_size + eligible_ids = ids[eligible_mask].numpy() + + if len(eligible_ids) < class_size: + raise ValueError( + f"Only found {len(eligible_ids)} identities with {sample_size}+ images." + ) + + return np.random.choice(eligible_ids, class_size, replace=False) + +def select_balanced_ids(dataset, class_size): + ids, counts = get_ids_and_counts(dataset=dataset) + sorted_indices = torch.argsort(counts, descending=True) + top_ids = ids[sorted_indices][:class_size].numpy() + return np.array(top_ids, dtype=int) + +def get_indices(dataset, identities, split_at): + train_indices = [] + test_indices = [] + + # We convert to numpy for faster searching with np.where + all_targets = np.array(dataset.targets) + + for person_id in identities: + # Get all indices for this specific person + indices = np.where(all_targets == person_id)[0] + + # Shuffle the indices for this person + np.random.shuffle(indices) + + # Split data based on your split_at value + train_indices.extend(indices[:split_at]) + test_indices.extend(indices[split_at:]) + + return train_indices, test_indices + + + +# optional function to get max amount of samples per class +def select_top_ids(dataset, class_size): + ids, counts = get_ids_and_counts(dataset=dataset) + + # sort by number of images (descending) + sorted_indices = torch.argsort(counts, descending=True) + + top_ids = ids[sorted_indices][:class_size].numpy() + + return np.array(top_ids, dtype=int) + + +def get_forget_retain_loaders(dataset: Dataset, forget_class_idx: int, batch_size: int = 32) -> tuple[DataLoader, DataLoader]: + """ + Splits an IdentitySubset or standard Dataset into forget and retain sets + based on a remapped target class index. + """ + # 1. Safely extract targets whether it's a standard dataset or a Subset wrapper + if hasattr(dataset, 'targets'): + targets = dataset.targets + elif hasattr(dataset, 'identity'): # Raw CelebA support + targets = dataset.identity + else: + # If it's an IdentitySubset or standard Subset, extract mapped targets sequentially + # This guarantees we get the 0 -> (n-1) remapped labels + targets = [dataset[i][1] for i in range(len(dataset))] + + if not isinstance(targets, torch.Tensor): + targets = torch.tensor(targets) + + # 2. Generate mask indices local to this subset + forget_indices = torch.where(targets == forget_class_idx)[0].tolist() + retain_indices = torch.where(targets != forget_class_idx)[0].tolist() + + # 3. Create PyTorch Subsets + forget_subset = Subset(dataset, forget_indices) + retain_subset = Subset(dataset, retain_indices) + + # 4. Wrap into clean DataLoaders + forget_loader = DataLoader(forget_subset, batch_size=batch_size, shuffle=False) + retain_loader = DataLoader(retain_subset, batch_size=batch_size, shuffle=True) + + 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 diff --git a/sets/CasiaFace.py b/sets/CasiaFace.py new file mode 100644 index 0000000..51e7449 --- /dev/null +++ b/sets/CasiaFace.py @@ -0,0 +1,21 @@ +import os +from torchvision import datasets +from torch.utils.data import Dataset +import torch +from .Data import Data + +class CasiaSet(Data): + def __init__(self, resolution: int = 224, sample_size = 190): + super().__init__(resolution = resolution, sample_size = sample_size) + + def get_set(self) -> Data: + path_str = "./datasets/casia-set" + path = os.path.abspath(path_str) + + if not os.path.exists(path): + raise FileNotFoundError(f"Unpacked dataset missing at {self.final_path}. Run Extractor.py first!") + print(f"Loading unpacked CASIA dataset from: {self.final_path}") + set = datasets.ImageFolder(root=path, transform=None) + # we set the target here + self.target = torch.tensor(set.targets) + return set diff --git a/sets/CelebA.py b/sets/CelebA.py new file mode 100644 index 0000000..481dcc5 --- /dev/null +++ b/sets/CelebA.py @@ -0,0 +1,20 @@ +from torchvision import datasets +from torch.utils.data import Dataset +import torch +from .Data import Data + +class CelebA(Data): + def __init__(self, resolution: int = 224, sample_size = 30): + super().__init__(resolution, sample_size = sample_size) + + def get_set(self): + set = datasets.CelebA( + root = "../data", + split='all', + target_type='identity', + download=False, + transform=None + ) + # set the target first + self.target = set.identity + return set diff --git a/sets/Data.py b/sets/Data.py new file mode 100644 index 0000000..ca2ff7e --- /dev/null +++ b/sets/Data.py @@ -0,0 +1,239 @@ +from torchvision import datasets, transforms +from torch.utils.data import Dataset, DataLoader, Subset, ConcatDataset +import torch +import numpy as np +import os + +from enum import Enum, auto +class Set_Name(Enum): + CELEBA = auto() + CASIAFACES = auto() + +# train set transform +def train_transform(res): + return transforms.Compose([ + # ResNet expects 224 x 224 res + # Inception expects 299 x 299 + transforms.Resize((res, res)), + transforms.RandomHorizontalFlip(p=0.5), + transforms.ColorJitter( + brightness=0.2, + contrast=0.2, + saturation=0.1 + ), + transforms.ToTensor(), + # normalise to + transforms.Normalize( + mean=[0.485, 0.456, 0.406], + std=[0.229, 0.224, 0.225] + ) + ]) + +# test set transform +def test_transform(res): + return transforms.Compose([ + transforms.Resize((res, res)), + transforms.ToTensor(), + transforms.Normalize( + mean=[0.485, 0.456, 0.406], + std=[0.229, 0.224, 0.225] + ) + ]) + +# Load data with 'identity' as target and transform it +def get_set(set_name:Set_Name): + return fetch_celeb_a() if set_name == Set_Name.CELEBA else fetch_casia_faces() + +def fetch_celeb_a(): + return datasets.CelebA( + root='./data', + split='all', + target_type='identity', + download=True, + transform=None + ) + +def fetch_casia_faces(): + # location of the data (path relative to project root) + final_path = os.path.abspath("./data/casia-set") + + if not os.path.exists(final_path): + raise FileNotFoundError( + f"Unpacked dataset not found at {final_path}. " + "Please run Extractor.py first!" + ) + + print(f"Loading unpacked CASIA dataset from: {final_path}") + + return datasets.ImageFolder( + root=final_path, + transform=None + ) + + + +def get_ids_and_counts(dataset): + + target = get_target(dataset=dataset) + return torch.unique( + input = target, + return_counts = True + ) + + + +# filter selected identities from dataset +# How many classes, how many images per class +def select_ids( dataset, sample_size, class_size): + ids, counts = get_ids_and_counts(dataset = dataset) + eligible_mask = counts >= sample_size + eligible_ids = ids[eligible_mask].numpy() + + if len(eligible_ids) < class_size: + raise ValueError( + f"Only found {len(eligible_ids)} identities with {sample_size}+ images." + ) + + # Randomly select identities + return np.random.choice(eligible_ids, class_size, replace=False) + +# optional function to get max amount of samples per class +def select_top_ids(dataset, class_size): + ids, counts = get_ids_and_counts(dataset = dataset) + + # sort by number of images (descending) + sorted_indices = torch.argsort(counts, descending = True) + + top_ids = ids[sorted_indices][:class_size].numpy() + + return np.array(top_ids, dtype=int) + + +def get_target(dataset): + """ + Unified target extractor. + Instantly reads raw dataset arrays or safely scales down to unpack wrapped Subsets. + """ + if hasattr(dataset, 'identity'): + # celebA + targets = dataset.identity + elif hasattr(dataset, 'targets'): + # others + targets = dataset.targets + else: + # If it's an IdentitySubset or standard Subset, extract mapped targets sequentially + # This guarantees we get the 0 -> (n-1) remapped labels + targets = [dataset[i][1] for i in range(len(dataset))] + + if not isinstance(targets, torch.Tensor): + targets = torch.tensor(targets) + + return targets + + +# split class images to train and test set. +def get_indices(dataset, identities, split_at, size = 30): + + if split_at >= size: # debug safety + raise ValueError(f"Split point ({split_at}) must be less than total size ({size}).") + + train_indices = [] + test_indices = [] + + target = get_target(dataset=dataset) + + #training_sample = int(sample_size * training_ratio) + np.random.seed(42) + for person_id in identities: + # Get all indices for this specific person + indices = torch.where(target == person_id)[0].numpy() + + # Shuffle the indices for this person + np.random.shuffle(indices) + + # split data to testing and training + train_indices.extend(indices[:split_at]) + test_indices.extend(indices[split_at:size]) + + return train_indices, test_indices + + + +def get_unlearning_loaders(dataset: Dataset, forget_class_idx: int, batch_size: int = 32) -> tuple[DataLoader, DataLoader]: + """ + Splits an IdentitySubset or standard Dataset into forget and retain sets + based on a remapped target class index. + """ + # extract targets + targets = get_target(dataset=dataset) + + # mask indices local to this subset + forget_indices = torch.where(targets == forget_class_idx)[0].tolist() + retain_indices = torch.where(targets != forget_class_idx)[0].tolist() + + # PyTorch Subsets + forget_subset = Subset(dataset, forget_indices) + retain_subset = Subset(dataset, retain_indices) + + # DataLoaders + forget_loader = DataLoader(forget_subset, batch_size=batch_size, shuffle=False) + retain_loader = DataLoader(retain_subset, batch_size=batch_size, shuffle=True) + + print(f"[Data Split] Local Class {forget_class_idx}: {len(forget_subset)} samples | Remaining Classes: {len(retain_subset)} samples.") + + 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/sets/Data_OOP.py b/sets/Data_OOP.py new file mode 100644 index 0000000..25d6822 --- /dev/null +++ b/sets/Data_OOP.py @@ -0,0 +1,174 @@ +import torch +import numpy as np +from abc import ABC, abstractmethod +from torchvision import transforms, datasets +from torch.utils.data import Dataset, DataLoader, Subset + +class Data(ABC): + """ + Handles image pipelines, identity filtering, indexing, and unlearning splits. + """ + def __init__(self, res: int = 224, sample_size = 30, class_size = 20): + self.res = res + self.sample_size = sample_size + self.class_size = class_size + self.target = None # will have to be set in get_set() + + def train_transform(self): + return transforms.Compose([ + # ResNet expects 224 x 224 res + # Inception expects 299 x 299 + transforms.Resize((self.res, self.res)), + transforms.RandomHorizontalFlip(p=0.5), + transforms.ColorJitter( + brightness=0.2, + contrast=0.2, + saturation=0.1 + ), + transforms.ToTensor(), + transforms.Normalize( + mean=[0.485, 0.456, 0.406], + std=[0.229, 0.224, 0.225] + ) + ]) + + + def test_transform(self): + return transforms.Compose([ + transforms.Resize((self.res, self.res)), + transforms.ToTensor(), + transforms.Normalize( + mean=[0.485, 0.456, 0.406], + std=[0.229, 0.224, 0.225] + ) + ]) + + @abstractmethod + def get_set(self)-> datasets.TorchDataset: + """Loads and returns the raw underlying PyTorch Dataset instance.""" + pass + + def get_targets(self) -> torch.Tensor: + return self.target + + def get_ids_and_counts(self) -> tuple[torch.Tensor, torch.Tensor]: + if self.target is None: + raise ValueError ("This should be called after the 'target' variable has been set.") + return torch.unique( + self.target, + return_counts=True + ) + + def select_ids(self) -> np.ndarray: + ids, counts = self.get_ids_and_counts() + eligible_mask = counts >= self.sample_size + eligible_ids = ids[eligible_mask].numpy() + + if len(eligible_ids) < self.class_size: + raise ValueError( + f"Only found {len(eligible_ids)} identities with {self.sample_size}+ images." + ) + + return np.random.choice(eligible_ids, self.class_size, replace=False) + + # Function to get max amount of samples per class + def select_top_ids(self) -> np.ndarray: + + ids, counts = self.get_ids_and_counts() + # sort by number of images (descending) + sorted_indices = torch.argsort(counts, descending=True) + top_ids = ids[sorted_indices][:self.class_size].numpy() + return np.array(top_ids, dtype=int) + + def get_indices(self, identities: np.ndarray, split_at: int, max_size: int = None) -> tuple[list, list]: + '''train_indices = [] + test_indices = [] + max_size = self.sample_size if max_size is None else max_size + + # Pull raw target tensor array using concrete implementation rules + all_targets = np.array(self.get_targets().cpu()) + np.random.seed(42) + + for person_id in identities: + indices = np.where(all_targets == person_id)[0] + np.random.shuffle(indices) + + # Constrain total sample tracking size if requested (e.g. CelebA ceiling) + current_pool = indices[:max_size] if max_size else indices + + if split_at >= len(current_pool): + raise ValueError(f"Split point ({split_at}) exceeds slice size ({len(current_pool)}) for class {person_id}.") + + train_indices.extend(current_pool[:split_at]) + test_indices.extend(current_pool[split_at:]) + + return train_indices, test_indices''' + if split_at >= self.sample_size: # debug safety + raise ValueError(f"Split point ({split_at}) must be less than total size ({self.sample_size}).") + + train_indices = [] + test_indices = [] + + #training_sample = int(sample_size * training_ratio) + np.random.seed(42) + target = self.get_targets() + for person_id in identities: + # Get all indices for this specific person + indices = torch.where(target == person_id)[0].numpy() + + # Shuffle the indices for this person + np.random.shuffle(indices) + + # split data to testing and training + train_indices.extend(indices[:split_at]) + test_indices.extend(indices[split_at:self.sample_size]) + + return train_indices, test_indices + + @staticmethod + def get_unlearn_loaders( + dataset: Dataset, + forget_class_idx: int, + batch_size: int = 32 + ) -> tuple[DataLoader, DataLoader]: + + """Splits an IdentitySubset into forget/retain parts based on local class index.""" + if hasattr(dataset, 'targets'): + targets = dataset.targets + elif hasattr(dataset, 'identity'): + targets = dataset.identity + else: + targets = [dataset[i][1] for i in range(len(dataset))] + + if not isinstance(targets, torch.Tensor): + targets = torch.tensor(targets) + + forget_indices = torch.where(targets == forget_class_idx)[0].tolist() + retain_indices = torch.where(targets != forget_class_idx)[0].tolist() + + forget_subset = Subset(dataset, forget_indices) + retain_subset = Subset(dataset, retain_indices) + + forget_loader = DataLoader(forget_subset, batch_size=batch_size, shuffle=False) + retain_loader = DataLoader(retain_subset, batch_size=batch_size, shuffle=True) + + print(f"[Data Split] Local Class {forget_class_idx}: {len(forget_subset)} samples | Remaining Classes: {len(retain_subset)} samples.") + + return forget_loader, retain_loader + + + @staticmethod + def getDataSet(set:SetType, sample_size): + # some test + if set == SetType.CASIA: + from sets.CasiaFace import CasiaFace + return CasiaFace(sample_size = sample_size) + if set == SetType.CELEBA: + from sets.CelebA import CelebA + return CelebA(sample_size=sample_size) + + +from enum import Enum, auto +class SetType(Enum): + CASIA = auto() + CELEBA = auto() diff --git a/sets/Extractor.py b/sets/Extractor.py new file mode 100644 index 0000000..dafe426 --- /dev/null +++ b/sets/Extractor.py @@ -0,0 +1,131 @@ +import os +import struct +from tqdm import tqdm +from collections import Counter +import hashlib + +def get_top_identities_binary(rec_path, idx_path, top_n=51): + """ + Pass 1: Scans the actual BINARY HEADERS in the .rec file. + This is the only way to be 100% sure which image belongs to whom. + """ + identity_counts = Counter() + + with open(idx_path, 'r') as f: + offsets = [int(line.strip().split('\t')[1]) for line in f.readlines()] + + print("Pass 1: Scanning binary headers to count identities...") + with open(rec_path, 'rb') as f: + for offset in tqdm(offsets): + f.seek(offset) + header_bin = f.read(32) # Read enough for the header + if len(header_bin) < 32: continue + + # MXNet Header format: [Flag, Label (float), ID, ID] + # The label is at offset 12 (float32) + label = int(struct.unpack('f', header_bin[12:16])[0]) + identity_counts[label] += 1 + + top_stats = identity_counts.most_common(top_n) + top_labels = {label for label, count in top_stats} + + print(f"\nTop {top_n} Identities by Binary Label:") + for label, count in top_stats: + print(f"ID: {label:<10} | Count: {count:<10}") + + return top_labels + +def extract_selected_binary(rec_path, idx_path, output_dir, top_labels): + if not os.path.exists(output_dir): + os.makedirs(output_dir) + + with open(idx_path, 'r') as f: + offsets = [int(line.strip().split('\t')[1]) for line in f.readlines()] + + print(f"\nPass 2: Extracting verified images...") + + # NEW: Keep track of how many images we've saved for each ID + # to avoid overwriting files. + save_counters = {label: 0 for label in top_labels} + total_extracted = 0 + + with open(rec_path, 'rb') as f: + for offset in tqdm(offsets): + f.seek(offset) + header_bin = f.read(32) + if len(header_bin) < 32: continue + + label = int(struct.unpack('f', header_bin[12:16])[0]) + + if label not in top_labels: + continue + + # Read image content + _, length_flag = struct.unpack('II', header_bin[:8]) + content_length = length_flag & ((1 << 31) - 1) + content = f.read(content_length) + + img_start = content.find(b'\xff\xd8') + if img_start == -1: continue + + target_folder = os.path.join(output_dir, str(label)) + os.makedirs(target_folder, exist_ok=True) + + # Use the counter for this specific label + current_count = save_counters[label] + img_filename = f"{current_count}.jpg" + img_path = os.path.join(target_folder, img_filename) + if(current_count > 405): + continue + + with open(img_path, 'wb') as img_f: + img_f.write(content[img_start:]) + + save_counters[label] += 1 + total_extracted += 1 + + print(f"\nDone! Extracted {total_extracted} total images.") + + + +def remove_duplicates(root_dir): + hashes = {} # hash -> first_filepath + duplicates_removed = 0 + + # Walk through every identity folder + for subdir, dirs, files in os.walk(root_dir): + for filename in tqdm(files, desc=f"Checking {os.path.basename(subdir)}"): + filepath = os.path.join(subdir, filename) + + # Calculate MD5 hash of the file + with open(filepath, 'rb') as f: + file_hash = hashlib.md5(f.read()).hexdigest() + + if file_hash in hashes: + # We've seen this image before! + os.remove(filepath) + duplicates_removed += 1 + else: + hashes[file_hash] = filepath + + print(f"\nClean-up complete. Removed {duplicates_removed} duplicate images.") + + +''' +if __name__ == "__main__": + # Point this to your unpacked Top 50 folder + target_dir = "./datasets/casia-set" + remove_duplicates(target_dir) +''' +if __name__ == "__main__": + base_dir = os.path.dirname(os.path.abspath(__file__)) + 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) + + # Step 2: Extract + extract_selected_binary(REC, IDX, OUT, top_verified_labels) + \ No newline at end of file diff --git a/sets/IdentitySubset.py b/sets/IdentitySubset.py new file mode 100644 index 0000000..e0953af --- /dev/null +++ b/sets/IdentitySubset.py @@ -0,0 +1,34 @@ +import torch + +class IdentitySubset(torch.utils.data.Dataset): + def __init__(self, dataset, indices, id_mapping, transform=None): + """ + Args: + dataset: The base dataset (CelebA or ImageFolder). + indices: List of indices belonging to the selected identities. + id_mapping: Dictionary mapping {old_label: new_label_0_to_N}. + transform: Transformations to apply to the images. + """ + self.dataset = dataset + self.indices = indices + self.id_mapping = id_mapping + self.transform = transform + + def __getitem__(self, idx): + # Access the base dataset using the stored index + img, old_id = self.dataset[self.indices[idx]] + + # Apply transform if provided + if self.transform: + img = self.transform(img) + + # Handle Label Logic: + # CelebA returns a Tensor, ImageFolder returns an int. + # We convert to a standard Python int for the dictionary lookup. + clean_id = old_id.item() if torch.is_tensor(old_id) else old_id + + # Map the original identity to our new 0 -> N-1 range + return img, self.id_mapping[clean_id] + + def __len__(self): + return len(self.indices) \ No newline at end of file diff --git a/unlearning/CertifiedUnlearning.py b/unlearning/CertifiedUnlearning.py new file mode 100644 index 0000000..92a0820 --- /dev/null +++ b/unlearning/CertifiedUnlearning.py @@ -0,0 +1,344 @@ +import torch +import torch.nn as nn +from torch.utils.data import DataLoader, RandomSampler +from torch.autograd import grad +from unlearning.Strategy import Strategy + +from sets.Data import * + +# Single-Batch Certified Unlearning for DNNs + +class CertifiedUnlearning(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 = 50000.0, + s1: int = 2, s2: int = 350, 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_params(self, model: nn.Module, named): + """ + Safely collects named parameter tuples while skipping + InceptionV3 auxiliary layers and tracking gradients. + """ + inner_model = getattr(model, "model", model) + + # Check if the current architecture is an Inception variant + is_inception = inner_model.__class__.__name__.lower() == "inception3" + + params_list = [] + for name, p in inner_model.named_parameters(): + if p.requires_grad: + # Discard the disconnected auxiliary training branch weights + if is_inception and "AuxLogits" in name: + continue + # CRITICAL: Append as a tuple so it can be unpacked as (name, param) + params_list.append((name, p)) + + return params_list if named else [e[1] for e in params_list] + + ''' + def _compute_loss_gradient(self, model, loader, device: torch.device): + + model.eval() + criterion = nn.CrossEntropyLoss(reduction='sum') + params = self.get_params(model, False) # [p for name, p in model.named_parameters() if p.requires_grad and "AuxLogits" not in name] + + + grad_accumulator = [torch.zeros_like(p, device = device) for p in params] + total_samples = 0''' + + # Accumulate true data cross-entropy gradients + ''' + 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) + + mini_grads = list(grad(loss, params, retain_graph=False)) + for i in range(len(grad_accumulator)): + grad_accumulator[i] += mini_grads[i].cpu().detach() + + # Empirical data mean conversion + for i in range(len(grad_accumulator)): + grad_accumulator[i] /= total_samples + + # L2 weight regularization + l2_reg_term = 0.0 + for param in params: + if param.requires_grad: + l2_reg_term += torch.sum(param ** 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] + ''' + ''' + with torch.set_grad_enabled(True): + 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) + + mini_grads = grad(loss, params, retain_graph=False) + for i in range(len(grad_accumulator)): + grad_accumulator[i] += mini_grads[i] + + # Empirical data mean conversion + for i in range(len(grad_accumulator)): + grad_accumulator[i] /= total_samples + + # OPTIMIZATION 2: Analytical L2 Regularization Gradient instead of autograd + # d/dx (l2_reg * x^2) = 2 * l2_reg * x + for i, param in enumerate(params): + grad_accumulator[i] += 2 * self.l2_reg * param.detach() + + return grad_accumulator + + 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) + elemwise_products = sum(torch.sum(g_elem * v_elem) for g_elem, v_elem in zip(first_grads, v)) + return grad(elemwise_products, params, create_graph=False)''' + ''' + def _stochastic_newton_update(self, g, dataset, model, device): + model.eval() + criterion = nn.CrossEntropyLoss() + params = self.get_params(model, False) # [p for p in model.parameters() if p.requires_grad] + h_res = [torch.zeros_like(p) for p in g] + + # progress + total_steps = self.s1 * self.s2 + step_interval = max(1, total_steps // 100) + + global_step = 0 + current_pct = 0 + + sampler = RandomSampler(dataset, replacement=True, num_samples=self.unlearn_bs * self.s2) + res_loader = DataLoader(dataset, batch_size=self.unlearn_bs, sampler=sampler) + res_iter = iter(res_loader) + + for _ in range(self.s1): + h_estimate = [p.clone() for p in g] + sampler = RandomSampler(dataset, replacement=True, num_samples=self.unlearn_bs * self.s2) + res_loader = DataLoader(dataset, batch_size=self.unlearn_bs, sampler=sampler) + res_iter = iter(res_loader) + + for _ in range(self.s2): + + global_step += 1 + + if global_step % step_interval == 0 and current_pct < 100: + current_pct += 1 + print(f"\rProgress: {current_pct}% done", end="", flush=True) + + 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 = sum(p.pow(2).sum() for p in params) + 'for param in params: + #if param.requires_grad: + l2_reg_term += torch.sum(param ** 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)): + h_estimate[k].copy_(h_estimate[k] + g[k] - (h_s[k] / self.scale)) + #h_res[k] += h_estimate[k] / self.scale + #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 + with torch.no_grad(): + for k in range(len(params)): + h_res[k] += h_estimate[k] / self.scale + + return [p / self.s1 for p in h_res] + ''' + + def _compute_loss_gradient(self, model, loader, device: torch.device): + model.eval() + criterion = nn.CrossEntropyLoss(reduction='sum') + params = self.get_params(model, False) + + # OPTIMIZATION 1: Keep accumulator on GPU device directly + grad_accumulator = [torch.zeros_like(p, device=device) for p in params] + total_samples = 0 + + with torch.set_grad_enabled(True): + 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) + + mini_grads = grad(loss, params, retain_graph=False) + for i in range(len(grad_accumulator)): + grad_accumulator[i] += mini_grads[i] + + # Empirical data mean conversion + for i in range(len(grad_accumulator)): + grad_accumulator[i] /= total_samples + + # OPTIMIZATION 2: Analytical L2 Regularization Gradient instead of autograd + # d/dx (l2_reg * x^2) = 2 * l2_reg * x + for i, param in enumerate(params): + grad_accumulator[i] += 2 * self.l2_reg * param.detach() + + return grad_accumulator + + def _hvp(self, loss, params, v): + first_grads = grad(loss, params, retain_graph=True, create_graph=True) + elemwise_products = sum(torch.sum(g_elem * v_elem) for g_elem, v_elem in zip(first_grads, v)) + return grad(elemwise_products, params, create_graph=False) + + def _stochastic_newton_update(self, g, dataset, model, device): + model.eval() + criterion = nn.CrossEntropyLoss() + params = self.get_params(model, False) + h_res = [torch.zeros_like(p, device=device) for p in g] + + total_steps = self.s1 * self.s2 + step_interval = max(1, total_steps // 100) + + global_step = 0 + current_pct = 0 + + # Create DataLoader outside or use optimal sampling + sampler = RandomSampler(dataset, replacement=True, num_samples=self.unlearn_bs * self.s2 * self.s1) + res_loader = DataLoader(dataset, batch_size=self.unlearn_bs, sampler=sampler) + res_iter = iter(res_loader) + + for _ in range(self.s1): + h_estimate = [p.clone() for p in g] + + for _ in range(self.s2): + global_step += 1 + + 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) + + # OPTIMIZATION 3: Clean up graph creation for loss & L2 + outputs = model(data) + loss = criterion(outputs, target) + + l2_reg_term = sum(p.pow(2).sum() for p in params) + loss += (self.l2_reg + self.gamma) * l2_reg_term + + h_s = self._hvp(loss, params, h_estimate) + + # OPTIMIZATION 4: Avoid deprecated .data, use detach() and in-place ops + with torch.no_grad(): + for k in range(len(params)): + h_estimate[k].copy_(h_estimate[k] + g[k] - (h_s[k] / self.scale)) + + + # feed back on status + if global_step % step_interval == 0 and current_pct < 100: + current_pct += 1 + print(f"\rProgress: {current_pct}% done", end="", flush=True) + + with torch.no_grad(): + for k in range(len(params)): + h_res[k] += h_estimate[k] / self.scale + + return [p / self.s1 for p in h_res] + + + def _certify(self, model, device, delta, full_certification): + certification = "full " if full_certification else "partial" + print(f"Performing {certification} certification") + + delta_idx = 0 + + # named_parameters to monitor layer positions + for name, param in self.get_params(model, True): + if param.requires_grad: + noise = self.std * torch.randn(param.data.size(), device=device) + + if full_certification: + param.data.add_(delta[delta_idx] + noise) + else: + # option for applying certification only to last layers + # deprecated + if "layer4" in name or "fc" in name: + param.data.add_(delta[delta_idx] + noise) + else: + # Keep early low-level vision filters entirely pristine + pass + + # Move to the next calculated Hessian vector block only after a valid update step + delta_idx += 1 + + 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 Forget set...") + g = self._compute_loss_gradient(model, forget_loader, device) + + print(">> Estimating non-convex inverse Hessian trajectories via Taylor series...") + dataset = retain_loader.dataset + delta = self._stochastic_newton_update(g, dataset, model, device) + + print(">> Applying parameter removal adjustments (-delta)...") + model = self._certify( + model= model, + device = device, + delta = delta, + full_certification = True + ) + + print(">> Certified Unlearning process completed successfully.") + return model + + + # overriden function + def _split_data(self, dataset): + # Certified unlearning does require both forget and retain sets + # split horizontaly. one class to forget and the rest to retain + 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/LinearFiltration.py b/unlearning/LinearFiltration.py new file mode 100644 index 0000000..e5331b4 --- /dev/null +++ b/unlearning/LinearFiltration.py @@ -0,0 +1,168 @@ +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): + super().__init__(target_class_index=target_class_index) + self.A = None + + def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module: + model.eval() + # Freeze internal params + for param in model.parameters(): + param.requires_grad = False + + device = next(model.parameters()).device + + return self.normalise( + model=model, + retain_loader=retain_loader, + forget_loader=forget_loader, + device=device, + forget_index=self.target_class_index + ) + + + def _get_classifier(self, model: nn.Module) -> nn.Linear: + + inner_model = getattr(model, "model", model) + + # looking for standard naming conventions in named modules + for name, module in inner_model.named_modules(): + # Check if it's our target linear layer + if (name == "fc" or name == "classifier") and isinstance(module, nn.Linear): + return module + + # Handle models (like EfficientNet) where the classifier is a Sequential block + if name == "classifier" and isinstance(module, nn.Sequential): + for sub_module in reversed(list(module.children())): + if isinstance(sub_module, nn.Linear): + return sub_module + + # scan backwards for the last Linear layer + for module in reversed(list(inner_model.modules())): + if isinstance(module, nn.Linear): + return module + + raise RuntimeError(f"Could not locate a linear classification head for {model.__class__.__name__}") + + + + def _compute_A(self, model, num_classes, loader, device): + model.eval() + + + # Initialize tracking tensors + sums = torch.zeros(num_classes, num_classes, device=device) + counts = torch.zeros(num_classes, device=device) + + with torch.no_grad(): + for inputs, targets in loader: + inputs, targets = inputs.to(device), targets.to(device) + + # the logit predictions + outputs = model(inputs) + + # One-hot encode targets to act as a routing mask + one_hot = torch.nn.functional.one_hot(targets, num_classes=num_classes).float() + + # add + sums += torch.t(one_hot) @ outputs + + # Sum columns of one-hot to get counts per class in this batch + counts += one_hot.sum(dim=0) + + # means + counts_safe = counts.unsqueeze(1) + print(f"COUNTS IS >>>>>>>>> {counts_safe}") + self.A = torch.where( + counts_safe > 0, + sums / counts_safe, + torch.zeros_like(sums) + ) + + # 9 + def _compute_z(self, tensor, forget_index): + + K = tensor.shape[0] + + + pi_a_f = torch.zeros(tensor.shape[1], device=tensor.device) + + t_1 = pi_a_f + # row vector for the forgotten class + a_f = tensor[forget_index, :] + + 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)) * a_f[mask_a_f].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 + + + # Normalisation filtration + def normalise(self, model, retain_loader, forget_loader, device, forget_index): + clf = self._get_classifier(model) + W = clf.weight.data.clone() + num_classes = W.shape[0] + + # we combine the data so we can calculate the mean of prdictions + full_loader = _combine_set(retain_loader, forget_loader) + # 8 + # Computing A is the most resource intensive part of this algorithm + # and to optimise the process, we computr it only once and re-use it + # because mean of all prdictions is the same for all + if self.A is None: + self._compute_A( + model = model, + num_classes = num_classes, + loader = full_loader, + device = device + ) + + # 9 + Z = self._compute_z(tensor=self.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(self.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(self.A) + # 11 + W_Z = B_Z @ A_inv @ W + + # 12 + clf = self._get_classifier(model) + clf.weight.copy_(W_Z) + + 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 new file mode 100644 index 0000000..14c99a6 --- /dev/null +++ b/unlearning/Strategy.py @@ -0,0 +1,62 @@ + +import torch.nn as nn +import time +import os +from pathlib import Path +from torch.utils.data import DataLoader +import Util + +class Strategy: + """Abstract base class for unlearning algorithms with automated, strategy-specific logging.""" + + 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 + + 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: + log_file = Path(f"reports/{self.strategy_name}/{model.__class__.__name__}/time_metrics.txt") + Util._initialize_log_file(log_file=log_file) + """ + Wraps the unlearning execution with automated timing and strategy-specific logging. + DO NOT override this method in subclasses. Override _run instead. + """ + + + + retain_loader, forget_loader = self._split_data(dataset) + + # record start time to evaluate efficiency + start_time = time.perf_counter() + # Execute core unlearning logic + processed_model = self._run(model, forget_loader, retain_loader) + + end_time = time.perf_counter() + execution_time = end_time - start_time + + # Log to the strategy's specific file + Util.log_metric( + log_file=log_file, + execution_time=execution_time + ) + + print(f"[{self.strategy_name}] Completed in {execution_time:.6f} seconds. Saved to {log_file}") + + return processed_model + + def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module: + """Subclasses implement their core unlearning logic here.""" + 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/WF.py b/unlearning/WF.py new file mode 100644 index 0000000..e3776d3 --- /dev/null +++ b/unlearning/WF.py @@ -0,0 +1,107 @@ +import torch +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 WeightF(Strategy): + """ + Verbatim implementation of Poppi et al.'s WF-Net framework modified + for static, single-class unlearning extraction. + """ + 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 + + def _optimise_filter(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader, device): + num_classes = model.fc.out_features + wf_model = WF_Net(original_model=model, num_classes=num_classes).to(device) + + # Optimize only the specific alpha masks + optimizer = optim.Adam([wf_model.alpha], lr=self.lr) + criterion = nn.CrossEntropyLoss() # Default reduction is 'mean' + + 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: + r_inputs, r_labels = r_inputs.to(device), r_labels.to(device) + + try: + f_inputs, _ = next(forget_iter) + except StopIteration: + forget_iter = iter(forget_loader) + f_inputs, _ = next(forget_iter) + f_inputs = f_inputs.to(device) + + optimizer.zero_grad() + + # Forward Pass + 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 Loss (Mean over batch) + loss_r = criterion(outputs_r, r_labels) + + # Forget Loss (Corrected to Mean over batch) + temperature = 1.0 + logits_f_scaled = outputs_f / temperature + + # 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() + + total_loss = loss_r + (self.gamma * loss_f) + total_loss.backward() + optimizer.step() + + t_loss_r += loss_r.item() + t_loss_f += loss_f.item() + 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() + + 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.") + + original_weights = target_conv.weight.data.clone().detach() + out_channels = original_weights.shape[0] + + # Freeze global network layers + for p in model.parameters(): + p.requires_grad = False + + wf_model = self._optimise_filter( + model, + forget_loader=forget_loader, + retain_loader=retain_loader, + device=device, + ) + + # --- 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/WeightFiltration.py b/unlearning/WeightFiltration.py new file mode 100644 index 0000000..e549a31 --- /dev/null +++ b/unlearning/WeightFiltration.py @@ -0,0 +1,139 @@ +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import DataLoader, ConcatDataset, Subset +from unlearning.Strategy import Strategy +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): + def __init__(self, + target_class_index: int, + arch, + num_classes: int = 20, + epochs: int = 6, + lr: float = 100.0, + gamma: float = 0.01, + lambda_1 = 25 + + ): + super().__init__(target_class_index=target_class_index) + self.epochs = epochs + self.lr = lr + self.gamma = gamma + self.num_classes = num_classes + self.wf_model = None + self.lambda_1 = lambda_1 + self.arch = arch + + + 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, + arch=self.arch, + size=self.num_classes, + original_model=model, + target_class_index=self.target_class_index + ) + + wf_net = wf_model.get() + optimizer = optim.SGD([wf_net.alpha], lr=self.lr) + # Use reduction='none' so we can manipulate individual item losses + criterion_none = nn.CrossEntropyLoss(reduction='none') + criterion_mean = nn.CrossEntropyLoss() + + for epoch in range(self.epochs): + t_loss_r, t_loss_f = 0.0, 0.0 + steps = 0 + + # 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) + f_inputs, f_labels = f_inputs.to(device), f_labels.to(device) + + optimizer.zero_grad() + + # retain data paired with randomly selected rows of alpha to compute the retaining loss + random_offset = torch.randint(0, self.num_classes - 1, size=r_labels.shape, device=device) + gate_signals_r = torch.where(random_offset >= r_labels, random_offset + 1, random_offset) + + outputs_r = wf_net(r_inputs, target_class_indices=gate_signals_r) + loss_r = criterion_mean(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) + + # Calculate loss for every single item in the batch at once + per_item_forget_loss = criterion_none(outputs_f, f_labels) + + # Use a scatter/sum approach to get class-wise losses without a Python loop + # Create a mask of unique classes present in this batch + unique_classes, inverse_indices = torch.unique(f_labels, return_inverse=True) + classes_in_batch = unique_classes.size(0) + + if classes_in_batch > 0: + # average CE loss per class + class_loss_sums = torch.zeros(classes_in_batch, device=device) + class_loss_sums.scatter_add_(0, inverse_indices, per_item_forget_loss) + + class_counts = torch.zeros(classes_in_batch, device=device) + class_counts.scatter_add_(0, inverse_indices, torch.ones_like(per_item_forget_loss)) + + mean_class_ce_loss = class_loss_sums / class_counts + + # Poppi et al. suggest employing reciprocal of the forget loss + # to avoid shortcomings of negative gradient approach + loss_f = torch.mean(1.0 / (mean_class_ce_loss + 1e-6)) + else: + loss_f = torch.tensor(0.0, device=device) + + # Regularisation penalty + loss_reg = torch.sum(1.0 - torch.sigmoid(wf_net.alpha)) + + # Backpropagation + total_loss = loss_r + (self.lambda_1 * loss_f) + (self.gamma * loss_reg) + total_loss.backward() + optimizer.step() + + # Keep tracking stats + t_loss_r += loss_r.item() + t_loss_f += loss_f.item() + 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() + + 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: + 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 + + return self.wf_model + + def _split_data(self, dataset): + return vertical_split( + dataset= dataset, + batch_size=32, + num_classes=self.num_classes + ) +