cleaned up code
This commit is contained in:
507
Tune.py
507
Tune.py
@@ -1,268 +1,339 @@
|
||||
# Finetuning a selected model
|
||||
# on a selected dataset
|
||||
# using selected parameters
|
||||
|
||||
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
|
||||
#from Data import *
|
||||
# from datasets.Casia import *
|
||||
#from IdentitySubset import IdentitySubset
|
||||
import Util
|
||||
from sets.Data import *
|
||||
from sets.IdentitySubset import IdentitySubset
|
||||
# models
|
||||
from architectures.Model import Model, Architecture
|
||||
|
||||
from unlearning.CertifiedUnlearning import CertifiedUnlearning
|
||||
from unlearning.LinearFiltration import LinearFiltration
|
||||
from unlearning.CertifiedRemoval import CertifiedRemoval
|
||||
from unlearning.WeightFiltration import WeightFiltration
|
||||
from eval.UnlearningAttack import UnlearningAttack
|
||||
from unlearning.Retrain import Retrain
|
||||
|
||||
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
|
||||
# Global Hyperparameters
|
||||
CLASS_SIZE:int = 20
|
||||
BATCH_SIZE:int = 16
|
||||
SAMPLE_SIZE:int = 30
|
||||
TRAINING_SAMPLE:int = 27
|
||||
|
||||
# depends on model architecture
|
||||
# ResNet, DenseNet = 224
|
||||
# Inception = 299
|
||||
RESOLUTION = 224
|
||||
RESOLUTION:int = 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()
|
||||
# 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
|
||||
|
||||
|
||||
for i in range(0,1):#CLASS_SIZE):
|
||||
FORGET_CLASS_IDX = i
|
||||
# Create model using Factory
|
||||
# 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
|
||||
|
||||
model = None
|
||||
dataset = get_set(set_name=dataset_name)
|
||||
print(f"> {dataset.__class__.__name__} dataset loaded")
|
||||
|
||||
if FINETUNE:
|
||||
model = Model.create(
|
||||
arch = arch,
|
||||
device = device,
|
||||
size = CLASS_SIZE)
|
||||
# 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')
|
||||
|
||||
# we may need to load existing model or finetune
|
||||
model.train(
|
||||
epochs = EPOCHS,
|
||||
loader = train_loader,
|
||||
rate = LR_RATE)
|
||||
# Isolate sample index partitions
|
||||
train_indices, test_indices = get_indices(
|
||||
dataset=dataset,
|
||||
identities=selected_identities,
|
||||
split_at=TRAINING_SAMPLE,
|
||||
size=SAMPLE_SIZE
|
||||
)
|
||||
|
||||
# save.
|
||||
file_name = f"{arch.name.lower}_{dataset_name.name.lower()}"
|
||||
model.save(filename=arch.name.lower())
|
||||
# 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_set = IdentitySubset(
|
||||
dataset=dataset,
|
||||
indices=train_indices,
|
||||
id_mapping=id_map,
|
||||
transform=tr_transform
|
||||
)
|
||||
|
||||
# done tuning
|
||||
|
||||
# EVALUATE
|
||||
te_transform = test_transform(RESOLUTION)
|
||||
# Testing
|
||||
test_data = IdentitySubset(
|
||||
dataset = dataset,
|
||||
test_set = IdentitySubset(
|
||||
dataset=dataset,
|
||||
indices=test_indices,
|
||||
id_mapping=id_map,
|
||||
transform=te_transform)
|
||||
transform=te_transform
|
||||
)
|
||||
|
||||
test_loader = DataLoader(
|
||||
test_data,
|
||||
batch_size=BATCH_SIZE,
|
||||
shuffle=False)
|
||||
print(f"> Total training images: {len(train_set)}")
|
||||
print(f'> Constants : Classes = {CLASS_SIZE}, Batch = {BATCH_SIZE}')
|
||||
|
||||
print(f"Total test images for these {CLASS_SIZE} classes: {len(test_data)}")
|
||||
# Create the base target model instance
|
||||
base_model = Model.create(arch=ARCH, device=device, size=CLASS_SIZE)
|
||||
|
||||
# Evaluate
|
||||
return {
|
||||
"device": device,
|
||||
"train_set": train_set,
|
||||
"test_set": test_set,
|
||||
"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_set = env_dict["train_set"]
|
||||
test_set = env_dict["test_set"]
|
||||
|
||||
test_loader = DataLoader(test_set, batch_size=BATCH_SIZE, shuffle=False)
|
||||
|
||||
|
||||
train_loader = DataLoader(train_set, 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_set)}")
|
||||
|
||||
# Evaluate original base checkpoint performance
|
||||
current_mode = "Finetuned"
|
||||
if FINETUNE:
|
||||
|
||||
#current_mode = "Finetuned"
|
||||
accuracy, report_dict = model.evaluate(
|
||||
loader = test_loader,
|
||||
mode=current_mode
|
||||
)
|
||||
|
||||
|
||||
# evaluate finetuned model
|
||||
try:
|
||||
accuracy, report_dict = model.evaluate(loader=test_loader, mode=current_mode)
|
||||
Util._log_to_csv(
|
||||
arch=model.__class__.__name__,
|
||||
mode = current_mode,
|
||||
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 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)
|
||||
# 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
|
||||
)
|
||||
|
||||
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)
|
||||
# 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_set = env_dict["train_set"]
|
||||
test_set = env_dict["test_set"]
|
||||
|
||||
# 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
|
||||
|
||||
# Segment specific unlearning loaders using class index boundaries
|
||||
retain_train_loader , forget_train_loader= get_unlearning_loaders(
|
||||
dataset=train_set, forget_class_idx=forget_class_idx, batch_size=BATCH_SIZE
|
||||
)
|
||||
retain_test_loader, forget_test_loader = get_unlearning_loaders(
|
||||
dataset=test_set, 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
|
||||
)
|
||||
# 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_set)
|
||||
strategy_in_use = strategy.__class__.__name__
|
||||
|
||||
if isinstance(unlearned,nn.Module):
|
||||
reloaded.model = unlearned
|
||||
else:
|
||||
reloaded = unlearned
|
||||
|
||||
|
||||
#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
|
||||
|
||||
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:
|
||||
|
||||
test_loader = DataLoader(test_set, batch_size=BATCH_SIZE, shuffle=False)
|
||||
|
||||
suite_runner.run_complete_evaluation(
|
||||
framework_name=strategy_in_use,
|
||||
test_loader = test_loader,
|
||||
target_class=forget_class_idx,
|
||||
forget_train_loader=forget_train_loader,
|
||||
forget_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 = 10
|
||||
|
||||
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
|
||||
)
|
||||
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
|
||||
)
|
||||
# Normalisation Filtration
|
||||
linear_filtration = LinearFiltration(
|
||||
|
||||
# Unlearning
|
||||
# train loaders passed here
|
||||
strategy.apply(reloaded.model, forget_train_loader, retain_train_loader)
|
||||
# Performance Analysis
|
||||
strategy_in_use = strategy.__class__.__name__
|
||||
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
|
||||
)
|
||||
|
||||
# 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)
|
||||
retrain = Retrain(
|
||||
target_class_index = 0,
|
||||
arch = ARCH,
|
||||
size = CLASS_SIZE,
|
||||
lr = 0.0001,
|
||||
epochs = 14
|
||||
|
||||
Util._log_to_csv(
|
||||
arch=reloaded.__class__.__name__,
|
||||
mode = current_mode,
|
||||
accuracy=accuracy,
|
||||
report_dict=report_dict,
|
||||
strategy=strategy_in_use
|
||||
)
|
||||
)
|
||||
|
||||
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
|
||||
)
|
||||
|
||||
|
||||
# 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
|
||||
)
|
||||
except KeyboardInterrupt:
|
||||
print("\nprogram interrupted. Exit!")
|
||||
break
|
||||
|
||||
Reference in New Issue
Block a user