diff --git a/Tune_new.py b/Tune_new.py new file mode 100644 index 0000000..a221003 --- /dev/null +++ b/Tune_new.py @@ -0,0 +1,342 @@ +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 +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 = 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" + + # 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_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) + + # 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_data) + 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: + + suite_runner.run_complete_evaluation( + framework_name=strategy_in_use, + target_class=forget_class_idx, + forget_loader=forget_train_loader, + retain_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 = 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 + # 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 + ) + # just a single class run before running all remaining classes. + #print(">> Single check run complete. Verification successful!") + #break + #dist_attacker.run_adversarial_evaluation() + #dist_attacker.run_incremental_evaluation(current_class_step=i) + + #if suite_runner is not None: + #suite_runner.shutdown_hook() + + except KeyboardInterrupt: + print("\nprogram interrupted. Exit!") + break