cleaned up code

This commit is contained in:
2026-07-08 23:53:07 +02:00
parent 31f461342e
commit 04875a62e9
23 changed files with 1362 additions and 715 deletions

View File

@@ -39,6 +39,7 @@ class UnlearningAttack:
# JS Distance is the square root of JS Divergence
return np.mean(jensenshannon(np.array(probs1), np.array(probs2), axis=1))
def calculate_a_dist(self, latent1, latent2):
@@ -79,10 +80,29 @@ class UnlearningAttack:
np.array(all_losses).reshape(-1, 1)
])
def run_parameter_space_mia(self, unlearned_model, shadow_model, forget_loader, retain_test_loader, device, index):
def run_parameter_space_mia(
self,
unlearned_model,
shadow_model,
forget_train_loader,
forget_test_loader,
device,
index
):
X_shadow_mem = self._extract_attack_features(shadow_model, forget_loader, device, index)
X_shadow_non = self._extract_attack_features(shadow_model, retain_test_loader, device, index)
X_shadow_mem = self._extract_attack_features(
shadow_model,
forget_train_loader,
device,
index
)
X_shadow_non = self._extract_attack_features(
shadow_model,
forget_test_loader,
device,
index
)
# Train MIA Classifier
min_train = min(len(X_shadow_mem), len(X_shadow_non))
@@ -93,8 +113,8 @@ class UnlearningAttack:
attack_classifier.fit(X_train, y_train)
# Evaluate MIA
X_eval_mem = self._extract_attack_features(unlearned_model, forget_loader, device, index)
X_eval_non = self._extract_attack_features(unlearned_model, retain_test_loader, device, index)
X_eval_mem = self._extract_attack_features(unlearned_model, forget_train_loader, device, index)
X_eval_non = self._extract_attack_features(unlearned_model, forget_test_loader, device, index)
min_test = min(len(X_eval_mem), len(X_eval_non))
X_test = np.vstack([X_eval_mem[:min_test], X_eval_non[:min_test]])
@@ -138,7 +158,15 @@ class UnlearningAttack:
# evaluate similarity of outputs
return accuracy_score(label_test, adversary.predict(data_test))
def run_logit_space_lookalike_mia(self, filtered_model, naive_retrained, forget_loader, device, target_class):
def run_logit_space_lookalike_mia(
self,
filtered_model,
naive_retrained,
test_loader,
device,
target_class
):
filtered_model.eval()
naive_retrained.eval()
@@ -146,7 +174,7 @@ class UnlearningAttack:
naive_logits = []
with torch.no_grad():
for data, _ in forget_loader:
for data, _ in test_loader:
data = data.to(device)
if filtered_model.__class__.__name__ == "WF_Module":
@@ -169,8 +197,20 @@ class UnlearningAttack:
# so that the metric is between 0 and 1.
return 2.0 * np.abs(lookalike_accuracy - 0.5)
def run_complete_evaluation(self, framework_name, target_class, forget_loader, retain_test_loader, unlearned_instance, base_shadow_instance, device):
"""Orchestrates specific pipeline routing cleanly using cached constructor parameters."""
def run_complete_evaluation(
self,
framework_name,
target_class,
forget_train_loader,
forget_test_loader,
test_loader,
unlearned_instance,
base_shadow_instance,
device
):
# load from disk if saved model available
target_dir = os.path.join("reports", framework_name)
os.makedirs(target_dir, exist_ok=True)
current_log_file = os.path.join(target_dir, "attack_values.csv")
@@ -179,12 +219,12 @@ class UnlearningAttack:
with open(current_log_file, "w") as f:
f.write("target_class, parameter_mia_accuracy, latent_distance_tell, lookalike_accuracy, A-Dist, JS-Dist\n")
# 1. Parameter-Space MIA and Latent Distance
# Parameter-Space MIA
parameter_mia_acc = self.run_parameter_space_mia(
unlearned_model=unlearned_instance.model,
shadow_model=base_shadow_instance.model,
forget_loader=forget_loader,
retain_test_loader=retain_test_loader,
forget_train_loader=forget_train_loader,
forget_test_loader=forget_test_loader,
device=device,
index=target_class
)
@@ -206,17 +246,23 @@ class UnlearningAttack:
lookalike_acc = self.run_logit_space_lookalike_mia(
filtered_model=unlearned_instance.model,
naive_retrained=reference_model_torch,
forget_loader=forget_loader,
test_loader=test_loader,
device=device,
target_class=target_class
)
# Calculate JS-Dist
js_dist = self.calculate_js_dist(unlearned_instance.model, reference_model_torch, forget_loader, device, target_class)
js_dist = self.calculate_js_dist(
unlearned_instance.model,
reference_model_torch,
forget_train_loader,
device,
target_class
)
# Extract features
unlearned_features = self._extract_attack_features(unlearned_instance.model, forget_loader, device, target_class)
retrained_features = self._extract_attack_features(reference_model_torch, forget_loader, device, target_class)
unlearned_features = self._extract_attack_features(unlearned_instance.model, forget_train_loader, device, target_class)
retrained_features = self._extract_attack_features(reference_model_torch, forget_train_loader, device, target_class)
# Calculate A-Dist using these features
a_dist = self.calculate_a_dist(unlearned_features, retrained_features)
@@ -225,7 +271,7 @@ class UnlearningAttack:
print(f"[{framework_name}] Class {target_class} | Parameter MIA: {parameter_mia_acc:.4f} Lookalike: {lookalike_acc:.4f}" )
with open(current_log_file, "a") as f:
f.write(f"{target_class},{parameter_mia_acc:.6f},{0.00000},{lookalike_acc:.6f}, {a_dist:.6f}, {js_dist:.6f}\n")
f.write(f"{target_class},{parameter_mia_acc:.6f},{lookalike_acc:.6f}, {a_dist:.6f}, {js_dist:.6f}\n")
return {
"parameter_mia_accuracy": parameter_mia_acc,