diff --git a/Data.py b/Data.py index abab889..e64a440 100644 --- a/Data.py +++ b/Data.py @@ -80,6 +80,10 @@ def select_top_ids(dataset, class_size): # 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 = [] diff --git a/OOP.py b/OOP.py new file mode 100644 index 0000000..a7f69c6 --- /dev/null +++ b/OOP.py @@ -0,0 +1,48 @@ + +# This is from wikipedia pseudocode implementation of a single +# ThresholdLogic Unit. +# done to make me understand OOP the Python way + + +# no need for brackets if not inheriting +class ThresholdLogicUnit: + + # define members in init + def __init__(self, threshold, weights): + self.threshold = threshold + self.weights = weights + + + # If a function has to make use of member variables + # it has to have self as param + def fire(self,inputs): + tots = 0 + #for i in range(0,inputs.size()): + for val, weight in zip(inputs, self.weights): + if val: + tots+= weight + + return tots > self.threshold + + +def main(): + # data + weights = [0.5, -0.2, 0.8] + threshold = 1.0 + + # Instantiate the class + tlu = ThresholdLogicUnit(threshold, weights) + + # Test + test_inputs = [1, 1, 0] + result = tlu.fire(test_inputs) + + print(f"The unit fired: {result}") + +# The "Guard" +if __name__ == "__main__": + main() + + + + \ No newline at end of file diff --git a/Tune.py b/Tune.py index d968b44..789bc85 100644 --- a/Tune.py +++ b/Tune.py @@ -12,6 +12,8 @@ from datasets.UniversalIdentitySubset import UniversalIdentitySubset as Identity # models from architectures.Model import Model, Architecture +from unlearning import LinearFiltration, WeightFiltration, CertifiedRemoval + # numbre of classes CLASS_SIZE = 20 # batch @@ -145,3 +147,17 @@ print("Evaluating loaded") reloaded.evaluate( loader = test_loader ) + + +strategies_to_test = [ + LinearFiltration(target_class_idx=12), + WeightFiltration(target_class_idx=12), + CertifiedRemoval(target_class_idx=12) +] + +# Run the comparative benchmark seamlessly +execution_profiles = {} +for strategy in strategies_to_test: + # Each iteration clones weights back to fine-tuned state before running + runtime = my_model.unlearn(strategy, forget_loader, retain_loader) + execution_profiles[strategy.__class__.__name__] = runtime diff --git a/architectures/Model.py b/architectures/Model.py index 256b050..f5292e2 100644 --- a/architectures/Model.py +++ b/architectures/Model.py @@ -7,6 +7,7 @@ import time import numpy as np from sklearn.metrics import classification_report from pathlib import Path +from unlearning.Strategy import Strategy class Model(ABC): def __init__(self, device, size): @@ -92,6 +93,21 @@ class Model(ABC): 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.network, forget_loader, retain_loader) + + elapsed_time = time.time() - start_time + print(f"{strategy.__class__.__name__} completed in {elapsed_time:.4f} seconds.") + + return elapsed_time + # Using the factory patern here diff --git a/architectures/ResNet50.py b/architectures/ResNet50.py index 74a4df4..860271b 100644 --- a/architectures/ResNet50.py +++ b/architectures/ResNet50.py @@ -23,14 +23,14 @@ class ResNet50(Model): m = models.resnet50(weights=models.ResNet50_Weights.DEFAULT) # freez all layers - for param in m.parameters(): - param.requires_grad = False + #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 + #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/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/unlearning/CertifiedRemoval.py b/unlearning/CertifiedRemoval.py new file mode 100644 index 0000000..e69de29 diff --git a/unlearning/LinearFiltration.py b/unlearning/LinearFiltration.py new file mode 100644 index 0000000..abb8306 --- /dev/null +++ b/unlearning/LinearFiltration.py @@ -0,0 +1,29 @@ + +import torch +from Strategy import Strategy + +class NormalizingLinearFiltration(Strategy): + def __init__(self, target_class_idx): + self.target_class_idx = target_class_idx + + def apply(self, model, forget_loader, retain_loader): + model.eval() + # Freeze parameters structurally + for param in model.parameters(): + param.requires_grad = False + + with torch.no_grad(): + # we modify only classification head + # Shape: [num_classes, feature_dim] + W = model.fc.weight.data + + # Compute the normalization transformation projection matrix (A) + # (In your full code, calculate A here matching Baumhauer et al.'s equations) + num_classes = W.shape[0] + A = torch.eye(num_classes, device=W.device) + # Mask/blend target class index distribution configurations here... + A[self.target_class_idx, :] = 0.0 + + # 3. Direct weight matrix override: W_filtered = A * W + sanitized_W = torch.mm(A, W) + model.fc.weight.copy_(sanitized_W) \ No newline at end of file diff --git a/unlearning/Strategy.py b/unlearning/Strategy.py new file mode 100644 index 0000000..e69de29 diff --git a/unlearning/WeightFiltration.py b/unlearning/WeightFiltration.py new file mode 100644 index 0000000..e69de29