strategies tested

.gitignore

@@ -14,3 +14,9 @@ trained_models/
 __pycache__/
 *.py[cod]
 lib64
+/reports
+*.bin
+*.idx
+*.rec
+*.lst
+property

IdentitySubset.py

@@ -1,20 +0,0 @@
-
-import torch
-
-class IdentitySubset(torch.utils.data.Dataset):
-    def __init__(self, dataset, indices, id_mapping, transform=None):
-        self.dataset = dataset
-        self.indices = indices
-        self.id_mapping = id_mapping
-        self.transform = transform
-
-    def __getitem__(self, idx):
-        img, old_id = self.dataset[self.indices[idx]]
-
-        if self.transform:
-            img = self.transform(img)
-
-        return img, self.id_mapping[old_id.item()]
-    
-    def __len__(self):
-        return len(self.indices)

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.

Tune.py

@@ -5,14 +5,17 @@
 from torch.utils.data import DataLoader
 from sklearn.metrics import classification_report
 import SetUp
-from Data import *
+#from Data import *
-#from datasets.Casia import *
+# from datasets.Casia import *
-from IdentitySubset import IdentitySubset
+#from IdentitySubset import IdentitySubset
-#from datasets.UniversalIdentitySubset import UniversalIdentitySubset as 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
@@ -20,7 +23,7 @@ import Util
 # numbre of classes 
 CLASS_SIZE = 20
 # batch 
-BATCH_SIZE = 32
+BATCH_SIZE = 16
 
 # size of images per class trainset + testset
 # 30 works best, more than that and we dont have enough data
@@ -50,7 +53,11 @@ arch = Architecture.RESNET50
 
 # DATA PREPARATION
 # load data set and prepare
-dataset = get_set()
+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,
@@ -62,7 +69,7 @@ dataset = get_set()
 # that way repeated calls return the same classes
 selected_identities = select_top_ids(
     dataset=dataset,
-    class_size= CLASS_SIZE,
+    class_size=CLASS_SIZE
 )
 
 print(f'> Selected {CLASS_SIZE} random identity classes from CelebA dataset.')
@@ -72,7 +79,8 @@ print(f'> A class has {TRAINING_SMPLE} train and {SAMPLE_SIZE-TRAINING_SMPLE} te
 train_indices, test_indices = get_indices(
     dataset = dataset, 
     identities = selected_identities, 
-    split_at = TRAINING_SMPLE
+    split_at = TRAINING_SMPLE,
+    size= SAMPLE_SIZE
 )
 
 # helps map class id to index
@@ -80,7 +88,7 @@ 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(res = RESOLUTION)
+tr_transform = train_transform(RESOLUTION)
 train_data = IdentitySubset(
     dataset=dataset, 
     indices=train_indices, 
@@ -100,8 +108,10 @@ print(f'> Constants : Classes = {CLASS_SIZE}, Batch = {BATCH_SIZE}, epochs = {EP
 # cuda if exists (it does here)
 device = SetUp.get_device()
 
+
 for i in range(0,CLASS_SIZE):
-# Create model using Factory
+    FORGET_CLASS_IDX = i
+    # Create model using Factory
     model = Model.create(
         arch = arch, 
         device = device, 
@@ -136,45 +146,101 @@ for i in range(0,CLASS_SIZE):
     print(f"Total test images for these {CLASS_SIZE} classes: {len(test_data)}")
 
     # Evaluate
-    mode, accuracy, report_dict = model.evaluate(
+    current_mode = "Finetuned"
-                                    loader = test_loader,
+    accuracy, report_dict = model.evaluate(
-                                    mode="finetunned"
+        loader = test_loader,
-                                    )
+        mode=current_mode
-    Util._log_to_csv(model=reloaded, mode = "finetuned", accuracy=accuracy, report_dict=report_dict, strategy="base")
+    )
 
-    # test again
+    Util._log_to_csv(
-    reloaded = Model.create(
+        arch=model.__class__.__name__, 
-        arch=arch,
+        mode = current_mode, 
-        device = device,
+        accuracy=accuracy, 
-        size = CLASS_SIZE
+        report_dict=report_dict, 
-        )
+        strategy="base"
-    reloaded.load(arch = arch)
+    )
-    print("fine tunned model loaded")
-    # reloaded.evaluate(
-    #   loader = test_loader
-    #)
 
-    # Unlearning
+    # unlearning algorithms
-    FORGET_CLASS_IDX = i
+    linear_filtration = LinearFiltration(target_class_idx=FORGET_CLASS_IDX)
+    #filtration.apply(reloaded.model)
 
-    forget_test_loader, retain_test_loader = get_forget_retain_loaders(
+    weight_filtration = WeightFiltration(num_classes = CLASS_SIZE,target_class_idx=FORGET_CLASS_IDX)
-        dataset=test_data, 
+    #weight_filtration.apply(reloaded.model)
+
+    certified_removal = CertifiedRemoval(removal_bound=0.05, epsilon=0.5, l2_reg=0.1)
+    #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
     )
 
-    #retain_test_loader = DataLoader(retain_test_loader.dataset, batch_size=BATCH_SIZE, shuffle=False)
+    # to evaluate
-
+    forget_test_loader, retain_test_loader = get_unlearning_loaders(
-    # 3. Instantiate and apply the Linear Filtration rule
+            dataset=test_data, 
-    filtration = LinearFiltration(target_class_idx=FORGET_CLASS_IDX)
+            forget_class_idx=FORGET_CLASS_IDX, 
-    filtration.apply(reloaded.model)
+            batch_size=BATCH_SIZE
+        )
 
 
-    # 4. Final Performance Analysis
+    strategies = [linear_filtration, weight_filtration, certified_removal]
-    print("\n--- Performance on Retained Classes")
+    #strategies = [linear_filtration]
-    mode, accuracy, report_dict = reloaded.evaluate(loader=retain_test_loader, mode="retain")
+    for strategy in strategies:
-    Util._log_to_csv(model=reloaded, mode = "retain", accuracy=accuracy, report_dict=report_dict, strategy="linearFiltration")
+        # 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
+        #)
 
-    print("\n--- Performance on Forgotten Class")
+        # Unlearning
-    mode, accuracy, report_dict = reloaded.evaluate(loader=forget_test_loader,mode="forget")
+        # train loaders passed here
-    Util._log_to_csv(model=reloaded, mode = "forgotten", accuracy=accuracy, report_dict=report_dict, strategy="linearFiltration")
+        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
+        )

Util.py

@@ -1,13 +1,14 @@
 
 from pathlib import Path
-from architectures.Model import Model
+import time
+import os
 
-def _log_to_csv(model:Model, mode, accuracy, report_dict, strategy):
+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()
+        #arch_name = model.__class__.__name__.lower()
         save_dir = Path(f"reports/{strategy}")
         save_dir.mkdir(parents=True, exist_ok=True)
-        csv_path = save_dir / f"{arch_name}-{mode}.csv"
+        csv_path = save_dir / f"{arch}-{mode}.csv"
         
         file_exists = csv_path.exists()
 
@@ -31,4 +32,17 @@ def _log_to_csv(model:Model, mode, accuracy, report_dict, strategy):
                 f.write(",".join(headers) + "\n")
             f.write(",".join(row) + "\n")
             
-        print(f">> Direct CSV metrics appended to {csv_path}")
+        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")

architectures/Model.py

@@ -8,8 +8,11 @@ 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
 
 class Model(ABC):
+    # need to add a weight decay here
     def __init__(self, device, size):
         self.device = device
         self.size = size
@@ -21,8 +24,10 @@ class Model(ABC):
 
     def train(self, epochs, loader, rate):
         criterion = nn.CrossEntropyLoss()
-        optimizer = optim.Adam(filter(lambda p: p.requires_grad, self.model.parameters()), lr=rate)
+        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)
+
         print(f"Starting training on {self.device}...")
         start_time = time.time()
         self.model.train()
@@ -37,41 +42,21 @@ class Model(ABC):
                 loss.backward()
                 optimizer.step()
                 total_loss += loss.item()
-            
+            scheduler.step()
+        
             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")
 
 
-    def evaluate(self, loader):
-        self.model.eval()
-        all_preds, all_labels = [], []
-        print("\nEvaluating...")
-
-        
-        
-        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())
-
-        classes = sorted(list(set(all_labels)))
-
-        accuracy = 100 * (np.array(all_preds) == np.array(all_labels)).sum() / len(classes)
-        print(f"Test Accuracy: {accuracy:.2f}%")
-        print(classification_report(all_labels, all_preds, labels=classes, zero_division=0))
-
-
     def save(self, filename=None):
         
         save_dir = Path("trained_models")
         save_dir.mkdir(parents=True, exist_ok=True)
 
-        # Determine filename (Default to class name if not provided)
+        # Filename (Default to class name if not provided)
         if filename is None:
             filename = f"{self.__class__.__name__.lower()}.pth"
         
@@ -150,65 +135,7 @@ class Model(ABC):
         
         # 3. Delegate file tracking to isolated helper method
         #self._log_to_csv(mode, accuracy,report_dict)
-        return mode, accuracy, report_dict
+        return accuracy, report_dict
-
-
-    def _log_to_csv(self, mode, accuracy, report_dict):
-        """Handles directory structures, file setups, and distinct CSV column formatting."""
-        arch_name = self.__class__.__name__.lower()
-        save_dir = Path("reports")
-        save_dir.mkdir(parents=True, exist_ok=True)
-        csv_path = save_dir / f"{arch_name}-{mode}.csv"
-        
-        file_exists = csv_path.exists()
-        '''
-        # Structure payload and headers based on evaluation slice type
-        if mode == "forget":
-            headers = ["accuracy", "precision", "recall", "f1-score"]
-            target_cls_str = str(classes[0])
-            metrics = report_dict[target_cls_str]
-            row = [
-                f"{accuracy / 100.0:.4f}",
-                f"{metrics['precision']:.4f}",
-                f"{metrics['recall']:.4f}",
-                f"{metrics['f1-score']:.4f}"
-            ]
-        else:
-            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}"
-            ]'''
-        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}")
 
 
 

reports/LinearFiltration_metrics.txt

@@ -1,35 +0,0 @@
-execution_time_sec
-0.001269
-0.001227
-0.001298
-0.001281
-0.001178
-0.001392
-0.001391
-0.001338
-0.001162
-0.001355
-0.001361
-0.001241
-0.001210
-0.001152
-0.001358
-0.001250
-0.001467
-0.001248
-0.001411
-0.001470
-0.001241
-0.001366
-0.001206
-0.001339
-0.001268
-0.002847
-0.001245
-0.001299
-0.001222
-0.001274
-0.001351
-0.001401
-0.001286
-0.001214

reports/resnet50-finetunned.csv

@@ -1,22 +0,0 @@
-accuracy,macro_precision,macro_recall,macro_f1,weighted_precision,weighted_recall,weighted_f1
-0.9333,0.9500,0.9333,0.9264,0.9500,0.9333,0.9264
-0.9167,0.9425,0.9167,0.9111,0.9425,0.9167,0.9111
-0.9333,0.9500,0.9333,0.9264,0.9500,0.9333,0.9264
-0.8667,0.9050,0.8667,0.8625,0.9050,0.8667,0.8625
-0.9000,0.9208,0.9000,0.8976,0.9208,0.9000,0.8976
-0.9000,0.9208,0.9000,0.8926,0.9208,0.9000,0.8926
-0.9667,0.9750,0.9667,0.9657,0.9750,0.9667,0.9657
-0.9000,0.9308,0.9000,0.9012,0.9308,0.9000,0.9012
-0.9833,0.9875,0.9833,0.9829,0.9875,0.9833,0.9829
-0.9500,0.9625,0.9500,0.9486,0.9625,0.9500,0.9486
-0.8667,0.9008,0.8667,0.8551,0.9008,0.8667,0.8551
-0.9167,0.9375,0.9167,0.9093,0.9375,0.9167,0.9093
-0.9000,0.9250,0.9000,0.8921,0.9250,0.9000,0.8921
-0.9000,0.9333,0.9000,0.9024,0.9333,0.9000,0.9024
-0.9500,0.9625,0.9500,0.9486,0.9625,0.9500,0.9486
-0.9167,0.9333,0.9167,0.9148,0.9333,0.9167,0.9148
-0.9167,0.9375,0.9167,0.9143,0.9375,0.9167,0.9143
-0.9667,0.9750,0.9667,0.9657,0.9750,0.9667,0.9657
-0.9333,0.9500,0.9333,0.9314,0.9500,0.9333,0.9314
-0.9000,0.9350,0.9000,0.8957,0.9350,0.9000,0.8957
-0.9000,0.9350,0.9000,0.9007,0.9350,0.9000,0.9007

reports/resnet50-forget.csv

@@ -1,21 +0,0 @@
-accuracy,macro_precision,macro_recall,macro_f1,weighted_precision,weighted_recall,weighted_f1
-0.0000,0.0000,0.0000,0.0000,0.0000,0.0000,0.0000
-0.0000,0.0000,0.0000,0.0000,0.0000,0.0000,0.0000
-0.0000,0.0000,0.0000,0.0000,0.0000,0.0000,0.0000
-0.0000,0.0000,0.0000,0.0000,0.0000,0.0000,0.0000
-0.0000,0.0000,0.0000,0.0000,0.0000,0.0000,0.0000
-0.0000,0.0000,0.0000,0.0000,0.0000,0.0000,0.0000
-0.0000,0.0000,0.0000,0.0000,0.0000,0.0000,0.0000
-0.0000,0.0000,0.0000,0.0000,0.0000,0.0000,0.0000
-0.0000,0.0000,0.0000,0.0000,0.0000,0.0000,0.0000
-0.0000,0.0000,0.0000,0.0000,0.0000,0.0000,0.0000
-0.0000,0.0000,0.0000,0.0000,0.0000,0.0000,0.0000
-0.0000,0.0000,0.0000,0.0000,0.0000,0.0000,0.0000
-0.0000,0.0000,0.0000,0.0000,0.0000,0.0000,0.0000
-0.0000,0.0000,0.0000,0.0000,0.0000,0.0000,0.0000
-0.0000,0.0000,0.0000,0.0000,0.0000,0.0000,0.0000
-0.0000,0.0000,0.0000,0.0000,0.0000,0.0000,0.0000
-0.0000,0.0000,0.0000,0.0000,0.0000,0.0000,0.0000
-0.0000,0.0000,0.0000,0.0000,0.0000,0.0000,0.0000
-0.0000,0.0000,0.0000,0.0000,0.0000,0.0000,0.0000
-0.0000,0.0000,0.0000,0.0000,0.0000,0.0000,0.0000

reports/resnet50-retain.csv

@@ -1,21 +0,0 @@
-accuracy,macro_precision,macro_recall,macro_f1,weighted_precision,weighted_recall,weighted_f1,
-0.9474,0.9605,0.9474,0.9459,0.9605,0.9474,0.9459
-0.9123,0.9395,0.9123,0.9064,0.9395,0.9123,0.9064
-0.9298,0.9526,0.9298,0.9244,0.9526,0.9298,0.9244
-0.8596,0.9000,0.8596,0.8553,0.9000,0.8596,0.8553
-0.9123,0.9298,0.9123,0.9103,0.9298,0.9123,0.9103
-0.9123,0.9342,0.9123,0.9045,0.9342,0.9123,0.9045
-0.9649,0.9737,0.9649,0.9639,0.9737,0.9649,0.9639
-0.8947,0.9272,0.8947,0.8960,0.9272,0.8947,0.8960
-1.0000,1.0000,1.0000,1.0000,1.0000,1.0000,1.0000
-0.9649,0.9737,0.9649,0.9639,0.9737,0.9649,0.9639
-0.8772,0.9132,0.8772,0.8650,0.9132,0.8772,0.8650
-0.9123,0.9342,0.9123,0.9045,0.9342,0.9123,0.9045
-0.9123,0.9342,0.9123,0.9045,0.9342,0.9123,0.9045
-0.9298,0.9605,0.9298,0.9328,0.9605,0.9298,0.9328
-0.9649,0.9737,0.9649,0.9639,0.9737,0.9649,0.9639
-0.9298,0.9430,0.9298,0.9283,0.9430,0.9298,0.9283
-0.9298,0.9474,0.9298,0.9278,0.9474,0.9298,0.9278
-0.9649,0.9737,0.9649,0.9639,0.9737,0.9649,0.9639
-0.9298,0.9474,0.9298,0.9278,0.9474,0.9298,0.9278
-0.8947,0.9316,0.8947,0.8902,0.9316,0.8947,0.8902

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

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

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=True, 
+            transform=None
+        )
+        # set the target first
+        self.target = set.identity
+        return set

sets/Data.py

@@ -1,9 +1,13 @@
-from torchvision import datasets, transforms, models
+from torchvision import datasets, transforms
 from torch.utils.data import Dataset, DataLoader, Subset
 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):
@@ -37,7 +41,10 @@ def test_transform(res):
     ])
 
 # Load data with 'identity' as target and transform it
-def get_set():
+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', 
@@ -46,17 +53,39 @@ def get_set():
         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(
-        dataset.identity, 
+        input = target,
-        return_counts=True
+        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)
+    ids, counts = get_ids_and_counts(dataset = dataset)
     eligible_mask = counts >= sample_size
     eligible_ids = ids[eligible_mask].numpy()
 
@@ -65,21 +94,43 @@ def select_ids( dataset, sample_size, class_size):
             f"Only found {len(eligible_ids)} identities with {sample_size}+ images."
         )
 
-    # Randomly select 50 identities
+    # 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)
+    ids, counts = get_ids_and_counts(dataset = dataset)
 
     # sort by number of images (descending)
-    sorted_indices = torch.argsort(counts, descending=True)
+    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):
 
@@ -89,11 +140,13 @@ def get_indices(dataset, identities, split_at, size = 30):
     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(dataset.identity == person_id)[0].numpy()
+        indices = torch.where(target == person_id)[0].numpy()
 
         # Shuffle the indices for this person
         np.random.shuffle(indices)
@@ -106,33 +159,23 @@ def get_indices(dataset, identities, split_at, size = 30):
 
 
 
-def get_forget_retain_loaders(dataset: Dataset, forget_class_idx: int, batch_size: int = 32) -> tuple[DataLoader, DataLoader]:
+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.
     """
-    # 1. Safely extract targets whether it's a standard dataset or a Subset wrapper
+    # extract targets
-    if hasattr(dataset, 'targets'):
+    targets = get_target(dataset=dataset)
-        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):
+    # mask indices local to this subset
-        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
+    # PyTorch Subsets
     forget_subset = Subset(dataset, forget_indices)
     retain_subset = Subset(dataset, retain_indices)
     
-    # 4. Wrap into clean DataLoaders
+    # DataLoaders
     forget_loader = DataLoader(forget_subset, batch_size=batch_size, shuffle=False)
     retain_loader = DataLoader(retain_subset, batch_size=batch_size, shuffle=True)
     

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()

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 > 200):
+                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, 'casia', 'train.rec')
+    IDX = os.path.join(base_dir, 'casia', 'train.idx')
+    OUT = os.path.join(base_dir, '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)
+    

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)

unlearning/CertifiedRemoval.py

@@ -1,77 +1,153 @@
 import torch
-import numpy as np
-from scipy.optimize import minimize
-from .Strategy import Strategy
 import torch.nn as nn
+from torch.utils.data import DataLoader
+from unlearning.Strategy import Strategy
 
 class CertifiedRemoval(Strategy):
-    """Implements Certified Removal for machine unlearning."""
+    """
-
+    Implements Certified Removal (Guo et al.) adapted for deep architectures
-    def __init__(self, model, data, labels, removal_bound, epsilon):
+    like ResNet50 by isolating and updating the final classification layer.
+    """
+    def __init__(self, removal_bound: float, epsilon: float, l2_reg: float = 0.1):
         super().__init__()
-        self.model = model
+        self.removal_bound = removal_bound  # gamma in the paper
-        self.data = data
+        self.epsilon = epsilon              # Privacy budget
-        self.labels = labels
+        self.l2_reg = l2_reg                # Lambda regularization term
-        self.removal_bound = removal_bound
-        self.epsilon = epsilon
 
-    def _run(self, model: nn.Module) -> nn.Module:
+    def _get_features(self, backbone: nn.Module, loader: DataLoader, device: torch.device):
-        """Runs the certified removal algorithm."""
+        """Passes data through the frozen ResNet backbone to extract embedding features."""
-        # 1. Linear Model Creation
+        backbone.eval()
-        # This is a simplification for demonstration purposes.  In a real implementation,
+        all_features = []
-        # you'd use more sophisticated methods to learn the parameters of the
+        all_labels = []
-        # 'removal' model based on the example being removed.
+        
-
+        with torch.no_grad():
-        def linear_model(x):
+            for inputs, labels in loader:
-            return torch.dot(x, torch.tensor([1, 1])) # Simplified Linear Model
+                inputs = inputs.to(device)
-
+                # Pass through backbone to get the 2048-dimensional feature vector
-        # 2. Optimization for Parameter Adjustment
+                features = backbone(inputs)
-        # Optimize the parameter values to minimize the loss while staying within bounds.
+                all_features.append(features.cpu())
-        original_params = torch.tensor([0.0, 0.0])  # Initial parameters for linear model
+                all_labels.append(labels.cpu())
-
-        def objective_function(params):
-            new_model = linear_model #use same function as defined above
-            return torch.sum(((new_model(self.data[0]) - self.labels)**2))
                 
-        result = minimize(objective_function, original_params, method='L-BFGS-B', bounds=[(-self.removal_bound, self.removal_bound)], options={'maxiter': 100})
+        return torch.cat(all_features, dim=0), torch.cat(all_labels, dim=0)
 
-        if not result.success:
+    def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
-            print("Warning: Optimization failed!")
+        """
-            print(result.message)
+        Entry point expected by your Model.unlearn() architecture interface.
-            return model #Return original if optimization fails
+        Applies Certified Removal strictly to the final linear layer (model.fc).
+        """
+        device = next(model.parameters()).device
+        
+        # Isolate the final NN (Fully connected) layer from the model
+        linear_head = model.fc
+        # Temporarily turn the fc layer into a identity pass-through
+        model.fc = nn.Identity() 
+        
+        print(">> Extracting deep features from model backbone...")
+        retain_features, retain_labels = self._get_features(model, retain_loader, device)
+        forget_features, forget_labels = self._get_features(model, forget_loader, device)
+        
+        # Restore the linear head back
+        model.fc = linear_head
+        
+        # Extract weights from the classification layer
+        # w shape: [num_classes, 2048]
+        w = model.fc.weight.data.clone().cpu()
+        
+        # Compute the Exact Hessian Matrix over the remaining (retained) features
+        # Formula: H = (X^T * X) / N + lambda * I
+        # this will be done on CPU. requires more ram so we cant afford to do it on VRAM
+        # print(">> Computing exact Hessian matrix...")
+        N_retain = retain_features.size(0)
+        # X_T_X = torch.matmul(retain_features.t(), retain_features)
+        # reg_matrix = self.l2_reg * torch.eye(retain_features.size(1))
+        hessian = self._compute_hessian(retain_features=retain_features, retain_features_size = N_retain)
+        
+        # Compute the gradient of the loss with respect to the forgotten data
+        # print(">> Calculating forget set gradients...")
+        # num_classes = w.size(0)
+        # Pass features through linear layer weights to get logits
+        # logits_forget = torch.matmul(forget_features, w.t()) 
+        # Apply softmax to get true class probabilities
+        # preds_softmax = torch.softmax(logits_forget, dim=1)
+        
+        # forget_labels_one_hot = torch.nn.functional.one_hot(forget_labels, num_classes=num_classes).float()
+        
+        #preds_forget = torch.matmul(forget_features, w.t())
+        #error = preds_forget - forget_labels_one_hot
+        # error = preds_softmax - forget_labels_one_hot
+        # grad_forget shape: [num_classes, 2048]
+        grad_forget = self._compute_loss_gradient(
+            forget_labels=forget_labels,
+            forget_features=forget_features, 
+            model_weights=w)
+        #torch.matmul(error.t(), forget_features) / forget_features.size(0)
 
-        new_params = result.x
+        # Compute the Newton step update via solving: H * Delta_W^T = Grad_forget^T
-        # 3. New Model Creation
+        delta_w = self._compute_newton_step(
+            tensor = hessian,
+            gradient= grad_forget
+        )
+        # print(">> Solving Newton step via system optimization...")
+        # try:
+        #     delta_w_t = torch.linalg.solve(Hessian, grad_forget.t())
+        #     delta_w = delta_w_t.t()
+        # except RuntimeError:
+        #     print(">> Warning: Hessian matrix is singular. Falling back to pseudo-inverse.")
+        #     delta_w = torch.matmul(grad_forget, torch.linalg.pinv(Hessian).t())
 
-        new_model = lambda x: torch.dot(x, new_params)
+        # Apply the Certified Removal update rule: W_new = W + Delta_W
-        return new_model
+        new_w = w + delta_w
+        # Calibrate noise based on your epsilon budget 
+        # (Guo et al. use a perturbation based on the regularization lambda and epsilon)
+        sigma = 2.0 / (self.l2_reg * self.epsilon)
+        noise = torch.randn_like(new_w) * (sigma / N_retain)
+        new_w = new_w + noise
 
+        # Theoretical Guarantee verification
+        norm_delta = torch.norm(delta_w).item()
+        if norm_delta > self.removal_bound:
+            print(f"!! Warning: Removal budget exceeded! Norm: {norm_delta:.4f} > Bound: {self.removal_bound}")
+        else:
+            print(f">> Certificate valid. Norm: {norm_delta:.4f} <= Bound: {self.removal_bound}")
 
-if __name__ == '__main__':
+        # Push updated parameters back into the model instance in-place
-    # Example Usage - Synthetic Data for Demonstration
+        model.fc.weight.data = new_w.to(device)
-    np.random.seed(42)  # For reproducibility
+        
-    n_samples = 100
+        print(">> Certified Removal process completed successfully.")
-    X = np.random.randn(n_samples, 2)
+        return model
-    y = (X[:, 0] + X[:, 1] > 0).astype(int)
+    
 
-    # Create a simple linear model for demonstration
+    # computing the hessian matrix
-    model = nn.Linear(2, 1)  # Simple linear classifier - PyTorch Version
+    def _compute_hessian(self, retain_features, retain_features_size):
-    optimizer = torch.optim.SGD(model.parameters(), lr=0.01) # Optimizer for training the linear model
+        print(">> Computing exact Hessian matrix...")
+        # N_retain = retain_features.size(0)
+        X_T_X = torch.matmul(retain_features.t(), retain_features)
+        reg_matrix = self.l2_reg * torch.eye(retain_features.size(1))
+        return (X_T_X / retain_features_size) + reg_matrix
+    
 
-    # Train a Linear Model
+    def _compute_loss_gradient(self, forget_features, forget_labels, model_weights):
-    for _ in range(100): #training loop
+        print(">> Calculating forget set gradients...")
-        optimizer.zero_grad()
+        num_classes = model_weights.size(0)
-        predictions = model(X)
+        # Pass features through linear layer weights to get logits
-        loss = torch.sum((predictions - y)**2)
+        logits_forget = torch.matmul(forget_features, model_weights.t()) 
-        loss.backward()
+        # Apply softmax to get true class probabilities
-        optimizer.step()
+        preds_softmax = torch.softmax(logits_forget, dim=1)
+        
+        forget_labels_one_hot = torch.nn.functional.one_hot(forget_labels, num_classes=num_classes).float()
+        
 
-    # Define parameters for Certified Removal
+        error = preds_softmax - forget_labels_one_hot
-    removal_bound = 1.0
+        # grad_forget shape: [num_classes, 2048]
-    epsilon = 0.1
+        return torch.matmul(error.t(), forget_features) / forget_features.size(0)
+    
 
-    # Create the CertifiedRemoval object with the trained model, data and labels
+    def _compute_newton_step(self,tensor, gradient):
-    certified_removal_obj = CertifiedRemoval(model, X, y, removal_bound, epsilon)
+        print(">> Solving Newton step via system optimization...")
-
+        try:
-    # Run Certified Removal
+            delta_w_t = torch.linalg.solve(tensor, gradient.t())
-    new_model = certified_removal_obj.apply(model)
+            delta_w = delta_w_t.t()
+        except RuntimeError:
+            print(">> Warning: Hessian matrix is singular. Falling back to pseudo-inverse.")
+            delta_w = torch.matmul(gradient, torch.linalg.pinv(tensor).t())
+        return delta_w

unlearning/LastK_Certified.py

@@ -0,0 +1,125 @@
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+from unlearning.Strategy import Strategy
+
+class LastKCertifiedRemoval(Strategy):
+    """
+    Implements Certified Removal (Guo et al.) scaled up to the last K layers 
+    of a ResNet50 network by flattening sub-graph parameters into a convex sub-problem.
+    """
+    def __init__(self, removal_bound: float, epsilon: float, l2_reg: float = 0.1):
+        super().__init__()
+        self.removal_bound = removal_bound
+        self.epsilon = epsilon
+        self.l2_reg = l2_reg
+
+    def _split_model(self, model: nn.Module):
+        """
+        Splits ResNet50 into a frozen feature backbone and an active unlearning head.
+        Here, 'Last K Layers' includes layer4 and the fc classification head.
+        """
+        # Feature Backbone: Everything up to layer3
+        backbone = nn.Sequential(
+            model.conv1,
+            model.bn1,
+            model.relu,
+            model.maxpool,
+            model.layer1,
+            model.layer2,
+            model.layer3
+        )
+        
+        # Active Head: Layer4, global pooling, and the final linear layer
+        unlearning_head = nn.Sequential(
+            model.layer4,
+            model.avgpool,
+            nn.Flatten(1),
+            model.fc
+        )
+        
+        return backbone, unlearning_head
+
+    def _get_intermediate_features(self, backbone: nn.Module, loader: DataLoader, device: torch.device):
+        """Extracts features from the exit point of the frozen backbone (post-layer3)."""
+        backbone.eval()
+        all_features = []
+        all_labels = []
+        
+        with torch.no_grad():
+            for inputs, labels in loader:
+                inputs = inputs.to(device)
+                features = backbone(inputs)
+                all_features.append(features.cpu())
+                all_labels.append(labels.cpu())
+                
+        return torch.cat(all_features, dim=0), torch.cat(all_labels, dim=0)
+
+    def apply(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
+        """
+        Extracts intermediate features and updates the parameters of the last blocks
+        using the exact inverse-Hessian influence step.
+        """
+        device = next(model.parameters()).device
+        
+        # 1. Slice the ResNet graph structural components
+        backbone, unlearning_head = self._split_model(model)
+        
+        print(">> Extracting intermediate structural features from layer3 exit...")
+        retain_feats, retain_labels = self._get_intermediate_features(backbone, retain_loader, device)
+        forget_feats, forget_labels = self._get_intermediate_features(backbone, forget_loader, device)
+        
+        # 2. Flatten target weights from the active head into a 1D optimization tensor
+        # For simplicity and mathematical stability, we isolate the final layer's weights
+        # inside the active head for the exact Hessian tracking step
+        target_layer = unlearning_head[-1] # This points straight to model.fc
+        w = target_layer.weight.data.clone().cpu()
+        
+        # 3. Compute Exact Hessian over intermediate embeddings
+        # ResNet50's layer4 expands channels to 2048, creating a 2048x2048 matrix context
+        print(">> Computing exact sub-graph Hessian matrix...")
+        N_retain = retain_feats.size(0)
+        
+        # Pool the feature maps if they haven't been flattened yet by the head module
+        if len(retain_feats.shape) > 2:
+            retain_flat = torch.mean(retain_feats, dim=[2, 3])
+            forget_flat = torch.mean(forget_feats, dim=[2, 3])
+        else:
+            retain_flat = retain_feats
+            forget_flat = forget_feats
+
+        X_T_X = torch.matmul(retain_flat.t(), retain_flat)
+        reg_matrix = self.l2_reg * torch.eye(retain_flat.size(1))
+        Hessian = (X_T_X / N_retain) + reg_matrix
+        
+        # 4. Calculate gradients relative to the forgotten target features
+        print(">> Calculating forget set gradients...")
+        num_classes = w.size(0)
+        forget_labels_one_hot = torch.nn.functional.one_hot(forget_labels, num_classes=num_classes).float()
+        
+        preds_forget = torch.matmul(forget_flat, w.t())
+        error = preds_forget - forget_labels_one_hot
+        grad_forget = torch.matmul(error.t(), forget_flat) / forget_flat.size(0)
+
+        # 5. Apply Newton Step optimization update
+        print(">> Inverting optimization subspace via system solver...")
+        try:
+            delta_w_t = torch.linalg.solve(Hessian, grad_forget.t())
+            delta_w = delta_w_t.t()
+        except RuntimeError:
+            print(">> Warning: Subspace Hessian is singular. Using pseudo-inverse fallback.")
+            delta_w = torch.matmul(grad_forget, torch.linalg.pinv(Hessian).t())
+
+        # 6. Apply Weight Adjustment Bounds Check
+        new_w = w + delta_w
+        norm_delta = torch.norm(delta_w).item()
+        if norm_delta > self.removal_bound:
+            print(f"!! Warning: Removal budget exceeded! Norm: {norm_delta:.4f} > Bound: {self.removal_bound}")
+        else:
+            print(f">> Certificate valid. Subspace Norm: {norm_delta:.4f} <= Bound: {self.removal_bound}")
+
+        # 7. Write weights directly back into the live ResNet50 instance
+        model.fc.weight.data = new_w.to(device)
+        
+        print(">> Last K Layers Certified Removal complete.")
+        return model

unlearning/LinearFiltration.py

@@ -2,13 +2,14 @@
 import torch
 import torch.nn as nn
 from .Strategy import Strategy
+from torch.utils.data import DataLoader
 
 class LinearFiltration(Strategy):
     def __init__(self, target_class_idx: int):
-        super().__init__()  # Automatically configures 'NormalizingLinearFiltration_metrics.txt'
+        super().__init__()
         self.target_class_idx = target_class_idx
 
-    def _run(self, model: nn.Module) -> nn.Module:
+    def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
         model.eval()
         for param in model.parameters():
             param.requires_grad = False
@@ -20,29 +21,28 @@ class LinearFiltration(Strategy):
             A = self._calculate_filtration_matrix(num_classes, self.target_class_idx, W.device)
             sanitized_W = torch.mm(A, W)
             model.fc.weight.copy_(sanitized_W)
+            # Filter the bias (if the layer uses one)
+            if model.fc.bias is not None:
+                b = model.fc.bias.data.clone()
+                # b is a 1D tensor of shape (num_classes), 
+                # so we use torch.mv (matrix-vector multiplication) or unsqueeze it
+                sanitized_b = torch.mv(A, b)
+                model.fc.bias.copy_(sanitized_b)
             
         return model
 
-    '''@staticmethod
-    def _calculate_filtration_matrix(num_classes: int, forget_class: int, device: torch.device) -> torch.Tensor:
-        A = torch.eye(num_classes, device=device)
-        num_remaining_classes = num_classes - 1
-        for j in range(num_classes):
-            if j == forget_class:
-                A[forget_class, j] = 0.0
-            else:
-                A[forget_class, j] = 1.0 / num_remaining_classes
-        return A'''
     @staticmethod
     def _calculate_filtration_matrix(num_classes: int, forget_class: int, device: torch.device) -> torch.Tensor:
         A = torch.eye(num_classes, device=device)
         num_remaining = num_classes - 1
         
-        # The row of the forgotten class should average the inputs of all other classes
+        # The row of the forgotten class should average all other classes
         for j in range(num_classes):
             if j == forget_class:
+                # we zero the forget class
                 A[forget_class, j] = 0.0
             else:
+                # and we distribute the output to the remaining
                 A[forget_class, j] = 1.0 / num_remaining
                 
         return A

unlearning/Strategy.py

@@ -2,6 +2,9 @@
 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."""
@@ -9,21 +12,10 @@ class Strategy:
     def __init__(self):
         # Dynamically set file name based on the class name (e.g., 'NormalizingLinearFiltration.txt')
         self.strategy_name = self.__class__.__name__
-        self.log_file = f"reports/{self.strategy_name}/metrics.txt"
+        self.log_file = Path(f"reports/{self.strategy_name}/metrics.txt")
-        self._initialize_log_file()
+        Util._initialize_log_file(log_file= self.log_file)
 
-    def _initialize_log_file(self):
+    def apply(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
-        """Creates a unique log file for this strategy with a header if it doesn't exist."""
-        if not os.path.exists(self.log_file):
-            with open(self.log_file, "w") as f:
-                f.write("execution_time_sec\n")
-
-    def log_metric(self, execution_time: float):
-        """Appends the execution time to this strategy's specific file."""
-        with open(self.log_file, "a") as f:
-            f.write(f"{execution_time:.6f}\n")
-
-    def apply(self, model: nn.Module) -> nn.Module:
         """
         Wraps the unlearning execution with automated timing and strategy-specific logging.
         DO NOT override this method in subclasses. Override _run instead.
@@ -31,17 +23,21 @@ class Strategy:
         start_time = time.perf_counter()
         
         # Execute core unlearning logic
-        processed_model = self._run(model)
+        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
-        self.log_metric(execution_time)
+        Util.log_metric(
+            log_file=self.log_file, 
+            execution_time=execution_time
+        )
+
         print(f"[{self.strategy_name}] Completed in {execution_time:.6f} seconds. Saved to {self.log_file}")
         
         return processed_model
 
-    def _run(self, model: nn.Module) -> nn.Module:
+    def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
         """Subclasses implement their core unlearning logic here."""
         raise NotImplementedError

unlearning/WeightFiltration.py

@@ -0,0 +1,140 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torch.utils.data import DataLoader
+from unlearning.Strategy import Strategy
+
+class WeightFiltration(Strategy):
+    """
+    Implements Poppi et al.'s Weight Filtering framework for linear layers.
+    Uses a standard functional hook to guarantee native PyTorch autograd tracking.
+    """
+    def __init__(self, num_classes: int, target_class_idx: int, epochs: int = 10, lr: float = 0.2, gamma: float = 10.0):
+        super().__init__()
+        self.num_classes = num_classes
+        self.target_class_idx = target_class_idx
+        self.epochs = epochs
+        self.lr = lr
+        self.gamma = gamma 
+        self.alpha = None
+        self.hook_handle = None
+
+    
+    def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
+        device = next(model.parameters()).device
+        model.eval() 
+        
+        # Locate layer4 for dynamic optimization
+        target_layer = model.layer4 if hasattr(model, 'layer4') else None
+        fc_layer = model.fc if hasattr(model, 'fc') and isinstance(model.fc, nn.Linear) else None
+        
+        if target_layer is None or fc_layer is None:
+            raise AttributeError("Model does not have the required layers.")
+        
+        # Match alpha dimensions to the channels outputted by layer4
+        num_features = fc_layer.weight.shape[1] 
+        self.alpha = nn.Parameter(torch.ones(self.num_classes, num_features, device=device) * 1.5)
+
+        # Freeze everything except our channel mask
+        for p in model.parameters():
+            p.requires_grad = False
+        self.alpha.requires_grad = True
+
+        # Hook into layer4 dynamically to run the untraining optimization
+        self.hook_handle = target_layer.register_forward_hook(self._get_hook())
+        # optimise the filter to maintain accuracy on retain set
+        # and decrease accuracy on forget set
+        self._optimise_filter(model, forget_loader, retain_loader, device)
+
+        # Remove the runtime hook
+        self.hook_handle.remove()
+        
+        # Transfer the channel suppression permanently into model.fc
+        with torch.no_grad():
+            mask = torch.sigmoid(self.alpha[self.target_class_idx]) # Shape: (num_features,)
+            
+            # Suppress the channels ONLY for the target class row in fc
+            fc_layer.weight[self.target_class_idx].copy_(
+                fc_layer.weight[self.target_class_idx] * mask
+            )
+            print(f">> Baked deep channel filter into Class {self.target_class_idx} weights.")
+            
+        return model
+
+    def _get_hook(self):
+        """
+        Filters the internal feature map channels of layer4.
+        The mask scales the channels across the batch.
+        """
+        def functional_hook(module, layer_input, layer_output):
+            # layer_output shape: (batch, channels, height, width) -> e.g., (16, 2048, 7, 7)
+            # self.alpha shape: (num_classes, channels) -> e.g., (20, 2048)
+            
+            # Extract 1D mask for the target class: (channels,)
+            mask = torch.sigmoid(self.alpha[self.target_class_idx])
+            
+            # Reshape mask to (1, channels, 1, 1) so it broadcasts over batch, height, and width
+            mask = mask.view(1, -1, 1, 1)
+            
+            # Scale the internal feature maps before they move to the next layer
+            return layer_output * mask
+            
+        return functional_hook
+    
+    
+    def _optimise_filter(self, model, forget_loader, retain_loader, device):
+        optimizer = optim.Adam([self.alpha], lr=self.lr)
+        criterion = nn.CrossEntropyLoss()
+
+        print(f"[{self.__class__.__name__}] Unlearning Class {self.target_class_idx} with gamma={self.gamma}...")
+
+        # To optimise this loop we will watch improvements after each optimisation
+        temp_forget_loss = None
+        # this can be adjusted to optimise the best escape point
+        # it is the value we set to evaluate performance improvement after each itteration.
+        # if improvement is less than this, then we break itteration.
+        threshold = 0.05
+
+        for epoch in range(self.epochs):
+            forget_iter = iter(forget_loader)
+            t_loss_r, t_loss_f = 0.0, 0.0
+            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()
+                
+                # Compute Losses
+                # The hook handles the weight filtering smoothly behind the scenes
+                loss_r = criterion(model(r_inputs), r_labels)
+                loss_f = -torch.sum((torch.ones_like(model(f_inputs)) / self.num_classes) * torch.log_softmax(model(f_inputs), dim=-1))
+                
+                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
+            forget_loss = t_loss_f / steps
+            print(f"  Epoch {epoch+1}/{self.epochs} | Retain Loss: {t_loss_r/steps:.4f} | Forget Loss: {forget_loss:.4f}")
+
+            if temp_forget_loss is not None:
+
+                improvement = temp_forget_loss - forget_loss
+                # if optimisation reaches a point of diminishing returns (improvements is less than threshold)
+                # we break the loop
+                if improvement < threshold:
+                    break
+            # else we update the lasst recorded loss.
+            temp_forget_loss = forget_loss