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.CertifiedRemoval import CertifiedRemoval 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"] # testing valuse * * #--------------------------------------------------------------------------- # S1 50 5 5 5 5 5 # S2 1000 200 1000 500 200 300 # BS 5 5 5 5 5 5 # scale 2000 500 8000 5000 10000 8000 # std 0.00001 0.00001 0.00001 0.00001 0.00001 0.00001 # Initialize the strategy hyperparameters matching standard settings # increase s2, decrease scale ---sweet spot '''certified_removal = CertifiedRemoval( target_class_index=forget_class_idx, s1=4, s2=350, # 350 best unlearn_bs=5, scale=6000.0, # 6000 was good std=0.00001 )''' '''certified_removal = CertifiedUnlearning( target_class_index=0, l2_reg=0.0005, gamma=0.1, scale=7000.0, s1=2, s2=350, std=1e-5, unlearn_bs=2 )''' # Segment specific unlearning loaders using class index boundaries forget_train_loader, retain_train_loader = get_unlearning_loaders( dataset=train_data, forget_class_idx=forget_class_idx, batch_size=BATCH_SIZE ) 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) if evaluate: reloaded.evaluate( loader=retain_test_loader, mode="finetuned" ) print("fine tunned model loaded into evaluation sandbox") # Execute strategic parameter unlearning step strategy.apply(reloaded.model, forget_train_loader, retain_train_loader) strategy_in_use = strategy.__class__.__name__ # 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__": # 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) finetuning = True # Unlearning Iterations for i in range(0, 1): # strategies # #certified_removal = CertifiedRemoval( # target_class_index=i, # s1=4, # s2=350, # 350 best # unlearn_bs=5, # scale=6000.0, # 6000 was good # std=0.00009 # ) certified_unlearning = CertifiedUnlearning( target_class_index=i, l2_reg=0.000002, gamma=0.1, scale= 20000,# 16400.0, # took ages to reach this sweet spot s1=2, s2=300, std=0.00001, unlearn_bs=16 ) # works perfectly linear_filtration = LinearFiltration( target_class_index=i ) weight_filtration = WeightFiltration( target_class_index=i, epochs=3, lr=0.5, gamma=150 ) strategies = [ certified_unlearning, # weight_filtration, # linear_filtration ] print(f"\n>>> Executing Unlearning Framework for Target Identity Index: {i} <<<") for strategy in strategies: run_unlearning_and_strategy_eval( runtime_environment, forget_class_idx=i, strategy=strategy, evaluate= not finetuning )