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

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