255 lines
8.1 KiB
Python
255 lines
8.1 KiB
Python
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
|
|
RESOLUTION = 224
|
|
ARCH = Architecture.RESNET18
|
|
|
|
|
|
# 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=10):
|
|
|
|
|
|
"""
|
|
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
|
|
#print(f"Starting training on {env_dict['device']}...")
|
|
model.train(epochs=epochs, loader=train_loader, rate=lr_rate)
|
|
model.save(filename=ARCH.name.lower())
|
|
print(f"Model saved to trained_models/{ARCH.name.lower()}.pth")
|
|
|
|
print(f"Total test images for these {CLASS_SIZE} classes: {len(test_data)}")
|
|
|
|
# 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=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=reloaded.__class__.__name__,
|
|
mode=domain["mode"],
|
|
accuracy=accuracy,
|
|
report_dict=report_dict,
|
|
strategy=strategy_in_use
|
|
)
|
|
|
|
|
|
# entry
|
|
if __name__ == "__main__":
|
|
|
|
|
|
try:
|
|
# Run Data Infrastructure and Architecture Builder
|
|
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)
|
|
|
|
|
|
# strategies
|
|
certified_unlearning = CertifiedUnlearning(
|
|
target_class_index=0,
|
|
l2_reg=0.000002,
|
|
gamma=0.1,
|
|
scale= 16400.0,# 16400.0, # took ages to reach this sweet spot
|
|
s1=2,
|
|
s2=300,
|
|
std=0.00001,
|
|
unlearn_bs=8
|
|
)
|
|
|
|
# works perfectly
|
|
linear_filtration = LinearFiltration(
|
|
|
|
target_class_index=0
|
|
)
|
|
|
|
weight_filtration = WeightFiltration(
|
|
target_class_index=0,
|
|
epochs=6,
|
|
lr=150.0,
|
|
gamma=0.001
|
|
)
|
|
|
|
strategies = [
|
|
certified_unlearning,
|
|
weight_filtration,
|
|
linear_filtration
|
|
]
|
|
# Unlearning Iteration
|
|
for i in range(0, CLASS_SIZE):
|
|
|
|
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,
|
|
evaluate = not finetuning
|
|
)
|
|
|
|
except KeyboardInterrupt:
|
|
print("program interrupted. Exit!")
|