402 lines
14 KiB
Python
402 lines
14 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from torch.utils.data import DataLoader
|
|
from sklearn.metrics import classification_report
|
|
import copy
|
|
|
|
# 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
|
|
from eval.UnlearningAttack import UnlearningAttack
|
|
from unlearning.Retrain import Retrain
|
|
|
|
|
|
# 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"
|
|
|
|
# evaluate finetuned model
|
|
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}")
|
|
|
|
|
|
# saves evaluation metrics to log files
|
|
def log_metrics(evaluation_domains, reloaded, strategy_in_use):
|
|
|
|
# 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,
|
|
mode=domain["mode"],
|
|
accuracy=accuracy,
|
|
report_dict=report_dict,
|
|
strategy=strategy_in_use
|
|
)
|
|
|
|
# performs MIA and ZRF attack on models and logs the results
|
|
def run_unlearning_and_attack_eval(forget_train_loader, retain_test_loader, reloaded, strategy_in_use, suite_runner, device, forget_class):
|
|
"""
|
|
Performs adversarial vulnerability stress tests (MIA and ZRF) in-memory
|
|
on the freshly unlearned model instance without saving it to disk first.
|
|
"""
|
|
if suite_runner is None:
|
|
raise ValueError("An active initialized UnlearningAttackSuite instance must be supplied.")
|
|
|
|
print(f"\n>>> Initializing Threat Model Stress Testing Suite for: {strategy_in_use}")
|
|
|
|
# 1. Dynamically map the white-box feature extraction hook to the active inner model
|
|
suite_runner.register_model_hook(reloaded.model)
|
|
|
|
# 2. Fire the complete evaluation suite using the isolated data split subsets
|
|
results = suite_runner.run_complete_evaluation(
|
|
target_class=forget_class,
|
|
framework_name=strategy_in_use,
|
|
forget_loader=forget_train_loader, # Members split from the train data partition
|
|
retain_test_loader=retain_test_loader, # Clean non-members split from validation data
|
|
device=device
|
|
)
|
|
|
|
print(f" [Attack Complete] Logit MIA AUC: {results['logit_mia_auc']:.4f} | "
|
|
f"Internal MIA AUC: {results['internal_mia_auc']:.4f} | "
|
|
f"ZRF Score: {results['zrf_score']:.4f}")
|
|
|
|
|
|
# performs MIA and ZRF attack on models and logs the results
|
|
def run_shaddow_attack_eval(forget_train_loader, retain_test_loader, reloaded, strategy_in_use, suite_runner, device, forget_class):
|
|
"""
|
|
Performs adversarial vulnerability stress tests matching the localized
|
|
shadow architecture specifications laid out in thesis Section 5.5.
|
|
"""
|
|
if suite_runner is None:
|
|
raise ValueError("An active initialized UnlearningAttackSuite instance must be supplied.")
|
|
|
|
print(f"\n>>> Initializing Threat Model Stress Testing Suite for: {strategy_in_use}")
|
|
|
|
# Instantiate a clean copy of the baseline trained model to serve as the Shadow reference proxy
|
|
# (Since finetuning is done once, we read its parameters cleanly from disk)
|
|
base_shadow = Model.create(arch=ARCH, device=device, size=CLASS_SIZE)
|
|
base_shadow.load(arch=ARCH)
|
|
|
|
# Execute the updated conditional attack framework
|
|
results = suite_runner.run_complete_evaluation(
|
|
framework_name=strategy_in_use,
|
|
target_class=forget_class,
|
|
forget_loader=forget_train_loader,
|
|
retain_test_loader=retain_test_loader,
|
|
unlearned_instance=reloaded, # The unlearned candidate model
|
|
base_shadow_instance=base_shadow, # The shadow proxy architecture
|
|
device=device
|
|
)
|
|
|
|
print(f" [Attack Complete] Adversary Binary Classification Accuracy: {results['mia_accuracy']:.4f}")
|
|
|
|
|
|
|
|
# Unlearning and strategy eval
|
|
def run_unlearning_and_strategy_eval(env_dict, forget_class_idx, strategy, evaluate = False, suite_runner=None):
|
|
"""
|
|
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)
|
|
|
|
# Clean un-manipulated snapshot to serve as the Parameter-Space shadow proxy reference
|
|
shadow_model = copy.deepcopy(reloaded)
|
|
|
|
if evaluate:
|
|
reloaded.evaluate(
|
|
loader=retain_test_loader, mode="finetuned"
|
|
)
|
|
|
|
print("fine tunned model loaded into evaluation sandbox")
|
|
|
|
# Execute strategic parameter unlearning step
|
|
# we are using only training data to unlearn.
|
|
# Test data is never touched here.
|
|
unlearned = strategy.apply(reloaded.model, train_data)
|
|
strategy_in_use = strategy.__class__.__name__
|
|
|
|
if isinstance(unlearned,nn.Module):
|
|
reloaded.model = unlearned
|
|
else:
|
|
reloaded = unlearned
|
|
|
|
|
|
|
|
is_retrained = isinstance(strategy, Retrain)
|
|
|
|
if is_retrained:
|
|
os.makedirs("trained_models", exist_ok=True)
|
|
reloaded.save(filename=f"class_{forget_class_idx}_retrained.pth")
|
|
|
|
|
|
|
|
# here we add a condition conditional statement
|
|
if suite_runner is not None:
|
|
|
|
suite_runner.run_complete_evaluation(
|
|
framework_name=strategy_in_use,
|
|
target_class=forget_class_idx,
|
|
forget_loader=forget_train_loader,
|
|
retain_test_loader=forget_test_loader,
|
|
unlearned_instance=reloaded,
|
|
base_shadow_instance=shadow_model,
|
|
device=device
|
|
)
|
|
|
|
|
|
# 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)"}
|
|
]
|
|
log_metrics(evaluation_domains, reloaded, 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
|
|
# switch finetuning for tests on strategies only,
|
|
# to avoid finetunning every time we test a strategy
|
|
finetuning = False
|
|
run_finetuning_or_baseline_eval(runtime_environment, run_training = finetuning)
|
|
# scale 16400.0 for ResNet
|
|
scale = 20100
|
|
# batch 8 for resNet,
|
|
unlearning_batches = 16
|
|
# regularis
|
|
# strategies
|
|
# implementation of Certified Removal for DNNs
|
|
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=350, # 300
|
|
std=0.00001, # 0.00001
|
|
unlearn_bs=unlearning_batches # 8 32 8
|
|
)
|
|
|
|
# Normalisation Filtration
|
|
linear_filtration = LinearFiltration(
|
|
|
|
target_class_index=0,
|
|
num_classes=CLASS_SIZE
|
|
)
|
|
# WF-Net
|
|
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
|
|
)
|
|
|
|
retrain = Retrain(
|
|
target_class_index = 0,
|
|
arch = ARCH,
|
|
size = CLASS_SIZE,
|
|
lr = 0.0001,
|
|
epochs = 14
|
|
|
|
)
|
|
|
|
strategies = [
|
|
retrain,
|
|
linear_filtration,
|
|
weight_filtration,
|
|
certified_unlearning,
|
|
]
|
|
suite_runner = UnlearningAttack(arch=ARCH, class_size=CLASS_SIZE)
|
|
# Unlearning Iteration
|
|
for i in range(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,
|
|
suite_runner=suite_runner
|
|
)
|
|
# just a single class run before running all remaining classes.
|
|
#print(">> Single check run complete. Verification successful!")
|
|
#break
|
|
#dist_attacker.run_adversarial_evaluation()
|
|
#dist_attacker.run_incremental_evaluation(current_class_step=i)
|
|
|
|
if suite_runner is not None:
|
|
suite_runner.shutdown_hook()
|
|
|
|
except KeyboardInterrupt:
|
|
print("\nprogram interrupted. Exit!")
|
|
break
|