238 lines
10 KiB
Python
238 lines
10 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import numpy as np
|
|
import os
|
|
from torch.utils.data import DataLoader
|
|
from sklearn.linear_model import LogisticRegression
|
|
from sklearn.metrics import accuracy_score
|
|
from architectures.Model import Model
|
|
|
|
class UnlearningAttack:
|
|
def __init__(self, arch, class_size):
|
|
"""
|
|
Initializes the robust verification suite with universal architecture metadata.
|
|
Matches Section 5.5 of the thesis text by implementing distinct
|
|
Parameter-Space and Logit-Space (Look-alike) attack pipelines uniformly.
|
|
"""
|
|
self.arch = arch
|
|
self.class_size = class_size
|
|
|
|
self.hook = None
|
|
self.model = None
|
|
self._hook_features = []
|
|
self.criterion = nn.CrossEntropyLoss(reduction='none')
|
|
self.collecting = False
|
|
|
|
def _hook_fn(self, module, input, output):
|
|
if not self.collecting:
|
|
return
|
|
flattened_embeddings = torch.flatten(output, 1)
|
|
self._hook_features.append(flattened_embeddings.detach())
|
|
|
|
def register_model_hook(self, inner_model):
|
|
if hasattr(inner_model, "original_model"):
|
|
core_model = inner_model.original_model
|
|
else:
|
|
core_model = inner_model
|
|
|
|
if hasattr(core_model, 'avgpool'):
|
|
pool_layer = core_model.avgpool
|
|
else:
|
|
pool_layer = None
|
|
for name, module in core_model.named_modules():
|
|
if 'pool' in name:
|
|
pool_layer = module
|
|
break
|
|
|
|
if pool_layer is None:
|
|
raise AttributeError("The target model architecture lacks an 'avgpool' layer block.")
|
|
|
|
self.hook = pool_layer.register_forward_hook(self._hook_fn)
|
|
|
|
def shutdown_hook(self):
|
|
if hasattr(self, 'hook') and self.hook:
|
|
self.hook.remove()
|
|
self.hook = None
|
|
|
|
def _extract_attack_features(self, target_model, loader, device, target_class):
|
|
target_model.eval()
|
|
all_probs = []
|
|
all_entropies = []
|
|
all_losses = []
|
|
self._hook_features = []
|
|
self.collecting = True
|
|
with torch.no_grad():
|
|
for data, targets in loader:
|
|
data, targets = data.to(device), targets.to(device)
|
|
if target_model.__class__.__name__ == "WF_Module":
|
|
gate_signals = torch.full(
|
|
(data.size(0),),
|
|
target_class,
|
|
dtype=torch.long,
|
|
device=data.device
|
|
)
|
|
outputs = target_model(data, target_class_indices=gate_signals)
|
|
else:
|
|
outputs = target_model(data)
|
|
|
|
probs = torch.softmax(outputs, dim=1)
|
|
all_probs.extend(probs.cpu().numpy())
|
|
|
|
entropy = -torch.sum(probs * torch.log(probs + 1e-10), dim=1)
|
|
all_entropies.extend(entropy.cpu().numpy())
|
|
|
|
loss = self.criterion(outputs, targets)
|
|
all_losses.extend(loss.cpu().numpy())
|
|
|
|
self.collecting = False
|
|
|
|
X_probs = np.array(all_probs)
|
|
X_entropy = np.array(all_entropies).reshape(-1, 1)
|
|
X_loss = np.array(all_losses).reshape(-1, 1)
|
|
X_features = np.hstack([X_probs, X_entropy, X_loss])
|
|
|
|
if self._hook_features:
|
|
compiled_latent = torch.cat(self._hook_features, dim=0).cpu().numpy()
|
|
else:
|
|
compiled_latent = np.zeros((len(X_features), 512))
|
|
|
|
return X_features, compiled_latent
|
|
|
|
def run_parameter_space_mia(self, unlearned_model, shadow_model, forget_loader, retain_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)
|
|
|
|
min_train = min(len(X_shadow_mem), len(X_shadow_non))
|
|
X_train = np.vstack([X_shadow_mem[:min_train], X_shadow_non[:min_train]])
|
|
y_train = np.concatenate([np.ones(min_train), np.zeros(min_train)])
|
|
|
|
attack_classifier = LogisticRegression(max_iter=1000)
|
|
attack_classifier.fit(X_train, y_train)
|
|
|
|
X_eval_mem, latent_features = self._extract_attack_features(unlearned_model, forget_loader, device, index)
|
|
X_eval_non, retain_latent = self._extract_attack_features(unlearned_model, retain_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]])
|
|
y_test = np.concatenate([np.ones(min_test), np.zeros(min_test)])
|
|
|
|
predictions = attack_classifier.predict(X_test)
|
|
mia_accuracy = accuracy_score(y_test, predictions)
|
|
|
|
clean_centroid = np.mean(retain_latent, axis=0)
|
|
forget_distances = np.linalg.norm(latent_features - clean_centroid, axis=1)
|
|
|
|
return mia_accuracy, np.mean(forget_distances)
|
|
|
|
def run_logit_space_lookalike_mia(self, filtered_model, naive_retrained, forget_loader, device, target_class):
|
|
filtered_model.eval()
|
|
naive_retrained.eval()
|
|
|
|
filtered_logits = []
|
|
naive_logits = []
|
|
|
|
with torch.no_grad():
|
|
for data, _ in forget_loader:
|
|
data = data.to(device)
|
|
|
|
if filtered_model.__class__.__name__ == "WF_Module":
|
|
gate_signals = torch.full((data.size(0),), target_class, dtype=torch.long, device=data.device)
|
|
out_filtered = filtered_model(data, target_class_indices=gate_signals)
|
|
else:
|
|
out_filtered = filtered_model(data)
|
|
|
|
out_naive = naive_retrained(data)
|
|
|
|
filtered_logits.append(out_filtered)
|
|
naive_logits.append(out_naive)
|
|
|
|
# Concatenate everything
|
|
filtered = torch.cat(filtered_logits, dim=0).cpu().numpy()
|
|
naive = torch.cat(naive_logits, dim=0).cpu().numpy()
|
|
|
|
# Z-Score Normalisation
|
|
filtered = (filtered - np.mean(filtered, axis=-1, keepdims=True)) / (np.std(filtered, axis=-1, keepdims=True) + 1e-8)
|
|
naive = (naive - np.mean(naive, axis=-1, keepdims=True)) / (np.std(naive, axis=-1, keepdims=True) + 1e-8)
|
|
|
|
# shuffle indices
|
|
num_images = len(filtered)
|
|
image_indices = np.arange(num_images)
|
|
np.random.shuffle(image_indices)
|
|
|
|
# split to train and test set
|
|
split = int(num_images * 0.7)
|
|
train_img_idx, test_img_idx = image_indices[:split], image_indices[split:]
|
|
|
|
# create a balanced test and train set
|
|
data_train = np.vstack([filtered[train_img_idx], naive[train_img_idx]])
|
|
data_test = np.vstack([filtered[test_img_idx], naive[test_img_idx]])
|
|
|
|
# labels for attcker (1 from unlearned and 0s to retrained)
|
|
# we do this because every output retrained gives us is a result of unseen
|
|
# but unlearned has seen these data.
|
|
label_train = np.concatenate([np.ones(len(train_img_idx)), np.zeros(len(train_img_idx))])
|
|
# test set
|
|
label_test = np.concatenate([np.ones(len(test_img_idx)), np.zeros(len(test_img_idx))])
|
|
|
|
|
|
adversary = LogisticRegression(max_iter=1000)
|
|
adversary.fit(data_train, label_train)
|
|
# evaluate similarity of outputs
|
|
lookalike_accuracy = accuracy_score(label_test, adversary.predict(data_test))
|
|
# 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."""
|
|
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")
|
|
|
|
if not os.path.exists(current_log_file):
|
|
with open(current_log_file, "w") as f:
|
|
f.write("target_class,parameter_mia_accuracy,latent_distance_tell,lookalike_accuracy\n")
|
|
|
|
self.register_model_hook(unlearned_instance.model)
|
|
|
|
# 1. Parameter-Space MIA and Latent Distance
|
|
parameter_mia_acc, latent_dist = 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,
|
|
device=device,
|
|
index=target_class
|
|
)
|
|
|
|
# we load a retrained model to evaluate look_alike tests
|
|
ghost_checkpoint_path = f"trained_models/class_{target_class}_retrained.pth"
|
|
|
|
if os.path.exists(ghost_checkpoint_path):
|
|
# Safe clean wrapper boot utilizing internal instance state properties
|
|
ghost_model_instance = Model.create(arch=self.arch, device=device, size=self.class_size)
|
|
state_dict = torch.load(ghost_checkpoint_path, map_location=device, weights_only=True)
|
|
ghost_model_instance.model.load_state_dict(state_dict)
|
|
reference_model_torch = ghost_model_instance.model
|
|
else:
|
|
raise FileNotFoundError(
|
|
f"Retrained weights not found at: {ghost_checkpoint_path}. \nDid you forget to save models or are they saved with a different path?"
|
|
)
|
|
|
|
lookalike_acc = self.run_logit_space_lookalike_mia(
|
|
filtered_model=unlearned_instance.model,
|
|
naive_retrained=reference_model_torch,
|
|
forget_loader=forget_loader,
|
|
device=device,
|
|
target_class=target_class
|
|
)
|
|
|
|
print(f"[{framework_name}] Class {target_class} | Parameter MIA: {parameter_mia_acc:.4f} | Latent Dist: {latent_dist:.4f} | Lookalike: {lookalike_acc:.4f}")
|
|
|
|
with open(current_log_file, "a") as f:
|
|
f.write(f"{target_class},{parameter_mia_acc:.6f},{latent_dist:.6f},{lookalike_acc:.6f}\n")
|
|
|
|
return {
|
|
"parameter_mia_accuracy": parameter_mia_acc,
|
|
"latent_distance": latent_dist,
|
|
"lookalike_accuracy": lookalike_acc
|
|
} |