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 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:int = 20 BATCH_SIZE:int = 16 SAMPLE_SIZE:int = 30 TRAINING_SAMPLE:int = 27 # depends on model architecture # ResNet, DenseNet = 224 # Inception = 299 RESOLUTION:int = 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 def get_device(): if torch.cuda.is_available(): # clear cach to boost memory # for new round torch.cuda.empty_cache() return torch.device("cuda") else: return torch.device("cpu") # 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. """ # get Cuda or CPU. device = 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_set = IdentitySubset( dataset=dataset, indices=train_indices, id_mapping=id_map, transform=tr_transform ) te_transform = test_transform(RESOLUTION) test_set = IdentitySubset( dataset=dataset, indices=test_indices, id_mapping=id_map, transform=te_transform ) print(f"> Total training images: {len(train_set)}") 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_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" # 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 ) # 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"] # 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 ) # 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 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, strategy_in_use = strategy_in_use ) # 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 = 20 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 ) except KeyboardInterrupt: print("\nprogram interrupted. Exit!") break