cleaned up for submission
This commit is contained in:
@@ -7,7 +7,6 @@ from unlearning.Strategy import Strategy
|
||||
from sets.Data import *
|
||||
|
||||
# Single-Batch Certified Unlearning for DNNs
|
||||
|
||||
class CertifiedUnlearning(Strategy):
|
||||
"""
|
||||
Implements Certified Unlearning for non-convex DNNs (Zhang et al.).
|
||||
@@ -47,8 +46,7 @@ class CertifiedUnlearning(Strategy):
|
||||
params_list = []
|
||||
for name, p in inner_model.named_parameters():
|
||||
if p.requires_grad:
|
||||
|
||||
# Append as a tuple so it can be unpacked as (name, param)
|
||||
# so it can be unpacked as (name, param)
|
||||
params_list.append((name, p))
|
||||
|
||||
return params_list if named else [e[1] for e in params_list]
|
||||
@@ -133,7 +131,7 @@ class CertifiedUnlearning(Strategy):
|
||||
h_estimate[k].copy_(h_estimate[k] + g[k] - (h_s[k] / self.scale))
|
||||
|
||||
|
||||
# feed back on status
|
||||
# feed back because this took agaes and it was not clear if things were working
|
||||
if global_step % step_interval == 0 and current_pct < 100:
|
||||
current_pct += 1
|
||||
print(f"\rProgress: {current_pct}% done", end="", flush=True)
|
||||
@@ -167,7 +165,7 @@ class CertifiedUnlearning(Strategy):
|
||||
# Keep early low-level vision filters entirely pristine
|
||||
pass
|
||||
|
||||
# Move to the next calculated Hessian vector block only after a valid update step
|
||||
# to the next calculated Hessian vector block
|
||||
delta_idx += 1
|
||||
|
||||
return model
|
||||
|
||||
@@ -124,7 +124,6 @@ class LinearFiltration(Strategy):
|
||||
if self.A is None:
|
||||
self._compute_A(
|
||||
model = model,
|
||||
#num_classes = num_classes,
|
||||
loader = forget_loader,
|
||||
device = device
|
||||
)
|
||||
@@ -147,6 +146,9 @@ class LinearFiltration(Strategy):
|
||||
A_inv = torch.linalg.pinv(A_t, rcond=1e-2)
|
||||
# 11
|
||||
W_temp = B_z @ A_inv @ W
|
||||
# Linear (Normalisation)filtration is done here,
|
||||
# next steps are added for safe automated evaluations
|
||||
|
||||
# with one class removed, we have a head W_temp for 19 classes.
|
||||
# but we have loaders for 20 classes for evaluation. so ,
|
||||
# map K-1 W back to K x K.
|
||||
@@ -164,6 +166,7 @@ class LinearFiltration(Strategy):
|
||||
# neutralise forget bias if exists
|
||||
if clf.bias is not None:
|
||||
n_bias = clf.bias.data.clone()
|
||||
# push it down to negative
|
||||
n_bias[forget_index] = -1e9
|
||||
clf.bias.copy_(n_bias)
|
||||
|
||||
|
||||
@@ -42,7 +42,7 @@ class Retrain(Strategy):
|
||||
print(f"No naive model found for class {self.target_class_index} retraining a new one")
|
||||
|
||||
|
||||
print(f">> Triggering Exact Unlearning Baseline (Retraining {self.arch.name} from pristine state)...")
|
||||
print(f"Retraining {self.arch.name} from pristine state)...")
|
||||
inner_model = getattr(model, "model", model)
|
||||
if hasattr(inner_model, "fc"):
|
||||
total_classes = inner_model.fc.out_features
|
||||
|
||||
107
unlearning/WF.py
107
unlearning/WF.py
@@ -1,107 +0,0 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.optim as optim
|
||||
from torch.utils.data import DataLoader
|
||||
from unlearning.Strategy import Strategy
|
||||
from .wf.WF_Net import WF_Net
|
||||
|
||||
class WeightF(Strategy):
|
||||
"""
|
||||
Verbatim implementation of Poppi et al.'s WF-Net framework modified
|
||||
for static, single-class unlearning extraction.
|
||||
"""
|
||||
def __init__(self, target_class_index: int, epochs: int = 10, lr: float = 0.2, gamma: float = 10.0):
|
||||
super().__init__(target_class_index=target_class_index)
|
||||
self.epochs = epochs
|
||||
self.lr = lr
|
||||
self.gamma = gamma
|
||||
|
||||
def _optimise_filter(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader, device):
|
||||
num_classes = model.fc.out_features
|
||||
wf_model = WF_Net(original_model=model, num_classes=num_classes).to(device)
|
||||
|
||||
# Optimize only the specific alpha masks
|
||||
optimizer = optim.Adam([wf_model.alpha], lr=self.lr)
|
||||
criterion = nn.CrossEntropyLoss() # Default reduction is 'mean'
|
||||
|
||||
for epoch in range(self.epochs):
|
||||
forget_iter = iter(forget_loader)
|
||||
t_loss_r, t_loss_f = 0.0, 0.0
|
||||
steps = 0
|
||||
|
||||
for r_inputs, r_labels in retain_loader:
|
||||
r_inputs, r_labels = r_inputs.to(device), r_labels.to(device)
|
||||
|
||||
try:
|
||||
f_inputs, _ = next(forget_iter)
|
||||
except StopIteration:
|
||||
forget_iter = iter(forget_loader)
|
||||
f_inputs, _ = next(forget_iter)
|
||||
f_inputs = f_inputs.to(device)
|
||||
|
||||
optimizer.zero_grad()
|
||||
|
||||
# Forward Pass
|
||||
outputs_r = wf_model(r_inputs, target_unlearn_class=self.target_class_index)
|
||||
outputs_f = wf_model(f_inputs, target_unlearn_class=self.target_class_index)
|
||||
|
||||
# Retain Loss (Mean over batch)
|
||||
loss_r = criterion(outputs_r, r_labels)
|
||||
|
||||
# Forget Loss (Corrected to Mean over batch)
|
||||
temperature = 1.0
|
||||
logits_f_scaled = outputs_f / temperature
|
||||
|
||||
# Compute uniform target entropy per-sample, then average over the batch
|
||||
log_probs_f = torch.log_softmax(logits_f_scaled, dim=-1)
|
||||
uniform_target = torch.ones_like(logits_f_scaled) / num_classes
|
||||
loss_f = -torch.sum(uniform_target * log_probs_f, dim=-1).mean()
|
||||
|
||||
total_loss = loss_r + (self.gamma * loss_f)
|
||||
total_loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
t_loss_r += loss_r.item()
|
||||
t_loss_f += loss_f.item()
|
||||
steps += 1
|
||||
|
||||
print(f" Epoch {epoch+1}/{self.epochs} | Retain Loss: {t_loss_r/steps:.4f} | Forget Loss: {t_loss_f/steps:.4f}")
|
||||
return wf_model
|
||||
|
||||
def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
|
||||
device = next(model.parameters()).device
|
||||
model.eval()
|
||||
|
||||
if hasattr(model, 'layer4') and len(model.layer4) > 1:
|
||||
target_conv = model.layer4[1].conv2
|
||||
else:
|
||||
raise AttributeError("Model architecture does not match expected ResNet-18 structure.")
|
||||
|
||||
original_weights = target_conv.weight.data.clone().detach()
|
||||
out_channels = original_weights.shape[0]
|
||||
|
||||
# Freeze global network layers
|
||||
for p in model.parameters():
|
||||
p.requires_grad = False
|
||||
|
||||
wf_model = self._optimise_filter(
|
||||
model,
|
||||
forget_loader=forget_loader,
|
||||
retain_loader=retain_loader,
|
||||
device=device,
|
||||
)
|
||||
|
||||
# --- PERMANENT BAKING STEP ---
|
||||
with torch.no_grad():
|
||||
# Grab the alpha mask vector for the forgotten class and cast to 4D tensor shape
|
||||
final_mask = torch.sigmoid(wf_model.alpha[self.target_class_index]).view(-1, 1, 1, 1)
|
||||
|
||||
# Apply filter masking permanently back onto the base layer
|
||||
target_conv.weight.copy_(original_weights * final_mask)
|
||||
|
||||
# Unfreeze architecture parameters for evaluations downstream
|
||||
for p in model.parameters():
|
||||
p.requires_grad = True
|
||||
|
||||
print(f">> Permanently altered {out_channels} convolutional filters in layer4 via WF-Net.")
|
||||
return model
|
||||
Reference in New Issue
Block a user