facebook's implementation

Tune_new.py

@@ -0,0 +1,189 @@
+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
+
+# Global Hyperparameters
+CLASS_SIZE = 20
+BATCH_SIZE = 16
+SAMPLE_SIZE = 30
+TRAINING_SAMPLE = 27
+RESOLUTION = 224
+ARCH = Architecture.RESNET50
+
+
+# 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.CELEBA
+
+    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 CelebA 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)
+
+    # Optional training configuration switch
+    if run_training:
+        train_loader = DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True)
+        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):
+    """
+    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"]
+
+    # Initialize the strategy hyperparameters matching standard settings
+    certified_removal = CertifiedRemoval(
+        target_class_index=forget_class_idx, 
+        removal_bound=0.05, 
+        epsilon=0.5, 
+        l2_reg=15
+    )
+
+    # 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)
+    print("fine tunned model loaded into evaluation sandbox")
+
+    # Execute strategic parameter unlearning step
+    certified_removal.apply(reloaded.model, forget_train_loader, retain_train_loader)
+    strategy_in_use = certified_removal.__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 
+    # switch finetuning for tests on strategies only
+    run_finetuning_or_baseline_eval(runtime_environment, run_training=True)
+    
+    # Unlearning Iterations
+    for i in range(0, 1):
+        print(f"\n>>> Executing Unlearning Framework for Target Identity Index: {i} <<<")
+        run_unlearning_and_strategy_eval(runtime_environment, forget_class_idx=i)