optimised

This commit is contained in:
2026-07-01 21:05:01 +02:00
commit 434d3b8198
34 changed files with 3286 additions and 0 deletions

24
.gitignore vendored Normal file
View File

@@ -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

50
DataAnalyser.py Normal file
View File

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

40
Note.md Normal file
View File

@@ -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.

67
Predict.py Normal file
View File

@@ -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.")

85
ReadME.md Normal file
View File

@@ -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
```

14
SetUp.py Normal file
View File

@@ -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")

268
Tune.py Normal file
View File

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

285
Tune_new.py Normal file
View File

@@ -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

51
Util.py Normal file
View File

@@ -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")

View File

@@ -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

View File

@@ -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

View File

@@ -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

View File

@@ -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")

225
architectures/Model.py Normal file
View File

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

22
architectures/ResNet18.py Normal file
View File

@@ -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

22
architectures/ResNet34.py Normal file
View File

@@ -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

36
architectures/ResNet50.py Normal file
View File

@@ -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

View File

@@ -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

230
architectures/WFNet.py Normal file
View File

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

View File

@@ -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

6
dependencies.txt Normal file
View File

@@ -0,0 +1,6 @@
torch
torchvision
gdown
numpy
scikit-learn
kagglehub

View File

@@ -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

167
sets/Casia.py Normal file
View File

@@ -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

21
sets/CasiaFace.py Normal file
View File

@@ -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

20
sets/CelebA.py Normal file
View File

@@ -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

239
sets/Data.py Normal file
View File

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

174
sets/Data_OOP.py Normal file
View File

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

131
sets/Extractor.py Normal file
View File

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

34
sets/IdentitySubset.py Normal file
View File

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

View File

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

View File

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

62
unlearning/Strategy.py Normal file
View File

@@ -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

107
unlearning/WF.py Normal file
View File

@@ -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

View File

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