Tinsae/Finetuning
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

clean up

Browse Source
This commit is contained in:
Tinsae
2026-07-08 20:36:49 +02:00
parent 90e33f074d
commit 31f461342e
3 changed files with 70 additions and 147 deletions
Download Patch File Download Diff File Expand all files Collapse all files

69
architectures/WFNet.py
View File

@@ -7,6 +7,68 @@ import numpy as np
from sklearn.metrics import classification_report
from architectures.Model import Model
'''class WF_Module(nn.Module):
"""
Pure PyTorch Neural Network module graph.
Keeps parameter registration and autograd tracking separate from
the framework's high-level Model abstractions to prevent recursion collisions.
"""
def __init__(self, original_model: nn.Module, num_classes: int):
super().__init__()
self.original_model = original_model
# Target layer for weight filtering (layer4 block 1 conv2 or conv3 depending on arch)
last_layer = original_model.layer4[1]
# Some versions are limited to 2 convolutional layers
if hasattr(last_layer, "conv3"):
self.target_conv = last_layer.conv3
else:
self.target_conv = last_layer.conv2
# Completely freeze the original ResNet parameters
for param in self.parameters():
param.requires_grad = False
# Initialize the alpha parameter matrix (Rows = Classes, Cols = Channels)
out_channels = self.target_conv.weight.shape[0]
self.alpha = nn.Parameter(torch.full((num_classes, out_channels), 3.0))'''
'''
Poppi et_al's Single-shot multiclass unlearning.
This calculation happens only once to generate the mask. once the mask is generated,
Unlearning and remembering becomes a matter of switching gates on and off.'''
'''
def forward(self, x: torch.Tensor, target_class_indices: torch.Tensor) -> torch.Tensor:
# we linearly loop through layers 1 to 4[block 1] (for ResNet)
# for i in M_{|L|} do l <- l[i]
x = self.original_model.maxpool(self.original_model.relu(self.original_model.bn1(self.original_model.conv1(x))))
x = self.original_model.layer1(x)
x = self.original_model.layer2(x)
x = self.original_model.layer3(x)
x = self.original_model.layer4[0](x)
# The second block execute its internal transformations natively
# This handles conv1->conv2 (ResNet18) or conv1->conv2->conv3 (ResNet50) automatically!
# Xi+1 <- l(Xi, ˆwl)
x = self.original_model.layer4[1](x)
# Apply mask dynamically to the completed block feature map
# wl <- αl[Yunl] ⊙ ˆwl
batch_alpha = self.alpha[target_class_indices]
mask = torch.sigmoid(batch_alpha).view(x.size(0), -1, 1, 1)
x = x * mask
# Remaining standard head steps
x = self.original_model.avgpool(x)
x = torch.flatten(x, 1)
# so here we are returning the output logits
# the result of classification is then
# argmax(x)
return self.original_model.fc(x)
'''
class WF_Module(nn.Module):
def __init__(self, original_model: nn.Module, num_classes: int, arch_enum):
super().__init__()
@@ -38,28 +100,23 @@ class WF_Module(nn.Module):
case arch_enum.RESNET18 | arch_enum.RESNET34 | arch_enum.RESNET50 | arch_enum.WIDE_RESNET:
return model.layer4[-1]
# --- GOOGLENET ---
case arch_enum.GOOGLENET:
return model.inception5b
# --- INCEPTION V3 ---
case arch_enum.INCEPTION:
return model.Mixed_7c
# --- DENSENET 121 ---
case arch_enum.DENSENET121:
return model.features.norm5
# --- EFFICIENTNET ---
case arch_enum.EFFICIENTNET:
return model.features[-1]
# --- SHUFFLENET ---
case arch_enum.SHUFFLENET:
return model.conv5
case _:
# Robust Fallback Strategy
# Fallback Strategy
target = None
for module in model.modules():
if isinstance(module, nn.Conv2d):

144
eval/UnlearningAttack.py
View File

@@ -17,10 +17,6 @@ class UnlearningAttack:
"""
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
@@ -45,109 +41,14 @@ class UnlearningAttack:
return np.mean(jensenshannon(np.array(probs1), np.array(probs2), axis=1))
def calculate_a_dist(self, latent1, latent2):
"""Calculates formal A-Distance: 2 * (1 - 2 * epsilon)."""
'''combined = np.vstack([latent1, latent2])
mean = np.mean(combined, axis=0)
std = np.std(combined, axis=0) + 1e-8
l1 = (latent1 - mean) / std
l2 = (latent2 - mean) / std
# 2. Use the same balanced split and regularization (C=0.01)
# as the look-alike method to ensure stability.
X = np.vstack([l1, l2])
y = np.concatenate([np.ones(len(l1)), np.zeros(len(l2))])
# Shuffle and split
idx = np.arange(len(X)); np.random.shuffle(idx)
X, y = X[idx], y[idx]
split = int(len(X) * 0.7)
clf = LogisticRegression(solver='liblinear').fit(X[:split], y[:split])
epsilon = 1.0 - accuracy_score(y[split:], clf.predict(X[split:]))'''
accuracy_score = self._comput_adversarial_accuracy(latent1, latent2, axis=0)
epsilon = 1 - accuracy_score
return 2.0 * np.abs(0.5 - epsilon)
'''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):
'''
# cnahgng to black box.
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'''
target_model.eval()
all_probs, all_entropies, all_losses = [], [], []
@@ -179,30 +80,7 @@ class UnlearningAttack:
])
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)'''
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)
@@ -226,7 +104,7 @@ class UnlearningAttack:
mia_accuracy = accuracy_score(y_test, predictions)
# Note: Latent distance is removed as it's not a black-box metric
return mia_accuracy, 0.0
return mia_accuracy
def _comput_adversarial_accuracy(self, filtered, naive, axis=0):
@@ -301,10 +179,8 @@ 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")
#self.register_model_hook(unlearned_instance.model)
# 1. Parameter-Space MIA and Latent Distance
parameter_mia_acc, latent_dist = self.run_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,
@@ -335,30 +211,24 @@ class UnlearningAttack:
target_class=target_class
)
# 1. Calculate JS-Dist (Logit-space probability comparison)
# Calculate JS-Dist
js_dist = self.calculate_js_dist(unlearned_instance.model, reference_model_torch, forget_loader, device, target_class)
# 2. Extract latent features for A-Dist
# We need features from both Unlearned and Retrained model
#_, unlearned_latent = self._extract_attack_features(unlearned_instance.model, forget_loader, device, target_class)
#_, retrained_latent = self._extract_attack_features(reference_model_torch, forget_loader, device, target_class)
# Extract features (now just one returned object)
# 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)
# Calculate A-Dist using these features
a_dist = self.calculate_a_dist(unlearned_features, retrained_features)
# 3. Calculate A-Dist (Replacing latent_distance)
#a_dist = self.calculate_a_dist(unlearned_latent, retrained_latent)
print(f"[{framework_name}] Class {target_class} | Parameter MIA: {parameter_mia_acc:.4f} | Latent Dist: {latent_dist:.4f} | Lookalike: {lookalike_acc:.4f}" )
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},{latent_dist:.6f},{lookalike_acc:.6f}, {a_dist:.6f}, {js_dist:.6f}\n")
f.write(f"{target_class},{parameter_mia_acc:.6f},{0.00000},{lookalike_acc:.6f}, {a_dist:.6f}, {js_dist:.6f}\n")
return {
"parameter_mia_accuracy": parameter_mia_acc,
"latent_distance": latent_dist,
"latent_distance": a_dist,
"lookalike_accuracy": lookalike_acc
}

4
unlearning/LinearFiltration.py
View File

@@ -12,10 +12,6 @@ class LinearFiltration(Strategy):
def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
model.eval()
# Freeze internal params
#for param in model.parameters():
#param.requires_grad = False
device = next(model.parameters()).device
return self.normalise(