293 lines
12 KiB
Python
293 lines
12 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import numpy as np
|
|
import os
|
|
from scipy.spatial.distance import jensenshannon
|
|
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.criterion = nn.CrossEntropyLoss(reduction='none')
|
|
self.collecting = False
|
|
|
|
|
|
def calculate_js_dist(self, model1, model2, loader, device, target_class):
|
|
"""Calculates Jensen-Shannon Distance between output probability distributions."""
|
|
model1.eval(); model2.eval()
|
|
probs1, probs2 = [], []
|
|
with torch.no_grad():
|
|
for data, _ in loader:
|
|
data = data.to(device)
|
|
# Handle WF_Module specific gate signal if needed
|
|
if model1.__class__.__name__ == "WF_Module":
|
|
gate = torch.full((data.size(0),), target_class, device=device)
|
|
p1 = torch.softmax(model1(data, target_class_indices=gate), dim=1)
|
|
|
|
else: # its a normal nn.Module
|
|
p1 = torch.softmax(model1(data), dim=1)
|
|
p2 = torch.softmax(model2(data), dim=1)
|
|
probs1.extend(p1.cpu().numpy()); probs2.extend(p2.cpu().numpy())
|
|
|
|
# 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):
|
|
|
|
accuracy_score = self._comput_adversarial_accuracy(latent1, latent2, axis=0)
|
|
epsilon = 1 - accuracy_score
|
|
|
|
return 2.0 * np.abs(0.5 - epsilon)
|
|
|
|
|
|
|
|
def _extract_attack_features(self, target_model, loader, device, target_class):
|
|
|
|
target_model.eval()
|
|
all_probs, all_entropies, all_losses = [], [], []
|
|
|
|
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 = torch.full((data.size(0),), target_class, dtype=torch.long, device=device)
|
|
outputs = target_model(data, target_class_indices=gate)
|
|
else:
|
|
outputs = target_model(data)
|
|
|
|
probs = torch.softmax(outputs, dim=1)
|
|
all_probs.extend(probs.cpu().numpy())
|
|
|
|
# Entropy: -sum(p * log(p))
|
|
log_probs = torch.log(probs + 1e-10)
|
|
entropy = -torch.sum(probs * log_probs, dim=1)
|
|
all_entropies.extend(entropy.cpu().numpy())
|
|
|
|
loss = self.criterion(outputs, targets)
|
|
all_losses.extend(loss.cpu().numpy())
|
|
|
|
# Combine output-based features only
|
|
return np.hstack([
|
|
np.array(all_probs),
|
|
np.array(all_entropies).reshape(-1, 1),
|
|
np.array(all_losses).reshape(-1, 1)
|
|
])
|
|
|
|
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_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))
|
|
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)
|
|
|
|
# Evaluate MIA
|
|
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]])
|
|
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)
|
|
|
|
# Note: Latent distance is removed as it's not a black-box metric
|
|
return mia_accuracy
|
|
|
|
def _comput_adversarial_accuracy(self, filtered, naive, axis=0):
|
|
|
|
# Z-Score Normalisation
|
|
filtered = (filtered - np.mean(filtered, axis = axis, keepdims=True)) / (np.std(filtered, axis = axis, keepdims = True) + 1e-8)
|
|
naive = (naive - np.mean(naive, axis = axis, keepdims = True)) / (np.std(naive, axis = axis, 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
|
|
return accuracy_score(label_test, adversary.predict(data_test))
|
|
|
|
def run_logit_space_lookalike_mia(
|
|
self,
|
|
filtered_model,
|
|
naive_retrained,
|
|
test_loader,
|
|
device,
|
|
target_class,
|
|
strategy_in_use,
|
|
):
|
|
|
|
filtered_model.eval()
|
|
naive_retrained.eval()
|
|
|
|
filtered_logits = []
|
|
naive_logits = []
|
|
|
|
with torch.no_grad():
|
|
for data, _ in test_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)
|
|
|
|
if(strategy_in_use == "LinearFiltration"):
|
|
mask = torch.ones(out_filtered.shape[1], dtype = torch.bool, device = device)
|
|
mask[target_class] = False
|
|
# out
|
|
out_filtered = out_filtered[:,mask]
|
|
out_naive = out_naive[:, mask]
|
|
|
|
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()
|
|
|
|
# evaluate similarity of outputs
|
|
lookalike_accuracy = self._comput_adversarial_accuracy(filtered=filtered, naive=naive, axis = -1)
|
|
# 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_train_loader,
|
|
forget_test_loader,
|
|
test_loader,
|
|
unlearned_instance,
|
|
base_shadow_instance,
|
|
device,
|
|
strategy_in_use
|
|
):
|
|
|
|
|
|
# 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")
|
|
|
|
if not os.path.exists(current_log_file):
|
|
with open(current_log_file, "w") as f:
|
|
f.write("target_class, parameter_mia_accuracy, lookalike_accuracy, A-Dist, JS-Dist\n")
|
|
|
|
# Parameter-Space MIA
|
|
parameter_mia_acc = self.run_parameter_space_mia(
|
|
unlearned_model=unlearned_instance.model,
|
|
shadow_model=base_shadow_instance.model,
|
|
forget_train_loader=forget_train_loader,
|
|
forget_test_loader=forget_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,
|
|
test_loader=test_loader,
|
|
device=device,
|
|
target_class=target_class,
|
|
strategy_in_use = strategy_in_use
|
|
)
|
|
|
|
# Calculate JS-Dist
|
|
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_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)
|
|
|
|
|
|
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},{lookalike_acc:.6f}, {a_dist:.6f}, {js_dist:.6f}\n")
|
|
|
|
return {
|
|
"parameter_mia_accuracy": parameter_mia_acc,
|
|
"latent_distance": a_dist,
|
|
"lookalike_accuracy": lookalike_acc
|
|
} |