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.GOOGLENET # 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!")