optimised
This commit is contained in:
344
unlearning/CertifiedUnlearning.py
Normal file
344
unlearning/CertifiedUnlearning.py
Normal file
@@ -0,0 +1,344 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils.data import DataLoader, RandomSampler
|
||||
from torch.autograd import grad
|
||||
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.).
|
||||
Uses a modified, stabilized stochastic Newton step using Taylor-expansion
|
||||
HVP estimation across the entire parameter space, capped with calibrated noise.
|
||||
"""
|
||||
def __init__(self, target_class_index: int, l2_reg: float = 0.0005,
|
||||
gamma: float = 0.01, scale: float = 50000.0,
|
||||
s1: int = 2, s2: int = 350, std: float = 0.001, unlearn_bs: int = 2):
|
||||
super().__init__(target_class_index)
|
||||
self.l2_reg = l2_reg
|
||||
self.gamma = gamma
|
||||
self.scale = scale
|
||||
self.s1 = s1
|
||||
self.s2 = s2
|
||||
self.std = std
|
||||
self.unlearn_bs = unlearn_bs
|
||||
|
||||
|
||||
|
||||
def get_params(self, model: nn.Module, named):
|
||||
"""
|
||||
Safely collects named parameter tuples while skipping
|
||||
InceptionV3 auxiliary layers and tracking gradients.
|
||||
"""
|
||||
inner_model = getattr(model, "model", model)
|
||||
|
||||
# Check if the current architecture is an Inception variant
|
||||
is_inception = inner_model.__class__.__name__.lower() == "inception3"
|
||||
|
||||
params_list = []
|
||||
for name, p in inner_model.named_parameters():
|
||||
if p.requires_grad:
|
||||
# Discard the disconnected auxiliary training branch weights
|
||||
if is_inception and "AuxLogits" in name:
|
||||
continue
|
||||
# CRITICAL: Append as a tuple 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]
|
||||
|
||||
'''
|
||||
def _compute_loss_gradient(self, model, loader, device: torch.device):
|
||||
|
||||
model.eval()
|
||||
criterion = nn.CrossEntropyLoss(reduction='sum')
|
||||
params = self.get_params(model, False) # [p for name, p in model.named_parameters() if p.requires_grad and "AuxLogits" not in name]
|
||||
|
||||
|
||||
grad_accumulator = [torch.zeros_like(p, device = device) for p in params]
|
||||
total_samples = 0'''
|
||||
|
||||
# Accumulate true data cross-entropy gradients
|
||||
'''
|
||||
for data, targets in loader:
|
||||
total_samples += targets.shape[0]
|
||||
data, targets = data.to(device), targets.to(device)
|
||||
outputs = model(data)
|
||||
loss = criterion(outputs, targets)
|
||||
|
||||
mini_grads = list(grad(loss, params, retain_graph=False))
|
||||
for i in range(len(grad_accumulator)):
|
||||
grad_accumulator[i] += mini_grads[i].cpu().detach()
|
||||
|
||||
# Empirical data mean conversion
|
||||
for i in range(len(grad_accumulator)):
|
||||
grad_accumulator[i] /= total_samples
|
||||
|
||||
# L2 weight regularization
|
||||
l2_reg_term = 0.0
|
||||
for param in params:
|
||||
if param.requires_grad:
|
||||
l2_reg_term += torch.sum(param ** 2)
|
||||
|
||||
reg_grads = list(grad(self.l2_reg * l2_reg_term, params))
|
||||
for i in range(len(grad_accumulator)):
|
||||
grad_accumulator[i] += reg_grads[i].cpu().detach()
|
||||
|
||||
return [p.to(device) for p in grad_accumulator]
|
||||
'''
|
||||
'''
|
||||
with torch.set_grad_enabled(True):
|
||||
for data, targets in loader:
|
||||
total_samples += targets.shape[0]
|
||||
data, targets = data.to(device), targets.to(device)
|
||||
outputs = model(data)
|
||||
loss = criterion(outputs, targets)
|
||||
|
||||
mini_grads = grad(loss, params, retain_graph=False)
|
||||
for i in range(len(grad_accumulator)):
|
||||
grad_accumulator[i] += mini_grads[i]
|
||||
|
||||
# Empirical data mean conversion
|
||||
for i in range(len(grad_accumulator)):
|
||||
grad_accumulator[i] /= total_samples
|
||||
|
||||
# OPTIMIZATION 2: Analytical L2 Regularization Gradient instead of autograd
|
||||
# d/dx (l2_reg * x^2) = 2 * l2_reg * x
|
||||
for i, param in enumerate(params):
|
||||
grad_accumulator[i] += 2 * self.l2_reg * param.detach()
|
||||
|
||||
return grad_accumulator
|
||||
|
||||
def _hvp(self, loss, params, v):
|
||||
first_grads = grad(loss, params, retain_graph=True, create_graph=True)
|
||||
elemwise_products = 0
|
||||
'''
|
||||
'''
|
||||
for grad_elem, v_elem in zip(first_grads, v):
|
||||
elemwise_products += torch.sum(grad_elem * v_elem)
|
||||
elemwise_products = sum(torch.sum(g_elem * v_elem) for g_elem, v_elem in zip(first_grads, v))
|
||||
return grad(elemwise_products, params, create_graph=False)'''
|
||||
'''
|
||||
def _stochastic_newton_update(self, g, dataset, model, device):
|
||||
model.eval()
|
||||
criterion = nn.CrossEntropyLoss()
|
||||
params = self.get_params(model, False) # [p for p in model.parameters() if p.requires_grad]
|
||||
h_res = [torch.zeros_like(p) for p in g]
|
||||
|
||||
# progress
|
||||
total_steps = self.s1 * self.s2
|
||||
step_interval = max(1, total_steps // 100)
|
||||
|
||||
global_step = 0
|
||||
current_pct = 0
|
||||
|
||||
sampler = RandomSampler(dataset, replacement=True, num_samples=self.unlearn_bs * self.s2)
|
||||
res_loader = DataLoader(dataset, batch_size=self.unlearn_bs, sampler=sampler)
|
||||
res_iter = iter(res_loader)
|
||||
|
||||
for _ in range(self.s1):
|
||||
h_estimate = [p.clone() for p in g]
|
||||
sampler = RandomSampler(dataset, replacement=True, num_samples=self.unlearn_bs * self.s2)
|
||||
res_loader = DataLoader(dataset, batch_size=self.unlearn_bs, sampler=sampler)
|
||||
res_iter = iter(res_loader)
|
||||
|
||||
for _ in range(self.s2):
|
||||
|
||||
global_step += 1
|
||||
|
||||
if global_step % step_interval == 0 and current_pct < 100:
|
||||
current_pct += 1
|
||||
print(f"\rProgress: {current_pct}% done", end="", flush=True)
|
||||
|
||||
try:
|
||||
data, target = next(res_iter)
|
||||
except StopIteration:
|
||||
res_iter = iter(res_loader)
|
||||
data, target = next(res_iter)
|
||||
|
||||
data, target = data.to(device), target.to(device)
|
||||
|
||||
outputs = model(data)
|
||||
|
||||
loss = criterion(outputs, target)
|
||||
l2_reg_term = sum(p.pow(2).sum() for p in params)
|
||||
'for param in params:
|
||||
#if param.requires_grad:
|
||||
l2_reg_term += torch.sum(param ** 2)
|
||||
loss += (self.l2_reg + self.gamma) * l2_reg_term
|
||||
|
||||
h_s = self._hvp(loss, params, h_estimate)
|
||||
|
||||
with torch.no_grad():
|
||||
for k in range(len(params)):
|
||||
h_estimate[k].copy_(h_estimate[k] + g[k] - (h_s[k] / self.scale))
|
||||
#h_res[k] += h_estimate[k] / self.scale
|
||||
#next_estimate = h_estimate[k].data + g[k].data - (h_s[k].data / self.scale)
|
||||
#h_estimate[k] = next_estimate.clone()
|
||||
del h_s, loss, outputs
|
||||
|
||||
#for k in range(len(params)):
|
||||
# h_res[k] = h_res[k] + h_estimate[k] / self.scale
|
||||
with torch.no_grad():
|
||||
for k in range(len(params)):
|
||||
h_res[k] += h_estimate[k] / self.scale
|
||||
|
||||
return [p / self.s1 for p in h_res]
|
||||
'''
|
||||
|
||||
def _compute_loss_gradient(self, model, loader, device: torch.device):
|
||||
model.eval()
|
||||
criterion = nn.CrossEntropyLoss(reduction='sum')
|
||||
params = self.get_params(model, False)
|
||||
|
||||
# OPTIMIZATION 1: Keep accumulator on GPU device directly
|
||||
grad_accumulator = [torch.zeros_like(p, device=device) for p in params]
|
||||
total_samples = 0
|
||||
|
||||
with torch.set_grad_enabled(True):
|
||||
for data, targets in loader:
|
||||
total_samples += targets.shape[0]
|
||||
data, targets = data.to(device), targets.to(device)
|
||||
outputs = model(data)
|
||||
loss = criterion(outputs, targets)
|
||||
|
||||
mini_grads = grad(loss, params, retain_graph=False)
|
||||
for i in range(len(grad_accumulator)):
|
||||
grad_accumulator[i] += mini_grads[i]
|
||||
|
||||
# Empirical data mean conversion
|
||||
for i in range(len(grad_accumulator)):
|
||||
grad_accumulator[i] /= total_samples
|
||||
|
||||
# OPTIMIZATION 2: Analytical L2 Regularization Gradient instead of autograd
|
||||
# d/dx (l2_reg * x^2) = 2 * l2_reg * x
|
||||
for i, param in enumerate(params):
|
||||
grad_accumulator[i] += 2 * self.l2_reg * param.detach()
|
||||
|
||||
return grad_accumulator
|
||||
|
||||
def _hvp(self, loss, params, v):
|
||||
first_grads = grad(loss, params, retain_graph=True, create_graph=True)
|
||||
elemwise_products = sum(torch.sum(g_elem * v_elem) for g_elem, v_elem in zip(first_grads, v))
|
||||
return grad(elemwise_products, params, create_graph=False)
|
||||
|
||||
def _stochastic_newton_update(self, g, dataset, model, device):
|
||||
model.eval()
|
||||
criterion = nn.CrossEntropyLoss()
|
||||
params = self.get_params(model, False)
|
||||
h_res = [torch.zeros_like(p, device=device) for p in g]
|
||||
|
||||
total_steps = self.s1 * self.s2
|
||||
step_interval = max(1, total_steps // 100)
|
||||
|
||||
global_step = 0
|
||||
current_pct = 0
|
||||
|
||||
# Create DataLoader outside or use optimal sampling
|
||||
sampler = RandomSampler(dataset, replacement=True, num_samples=self.unlearn_bs * self.s2 * self.s1)
|
||||
res_loader = DataLoader(dataset, batch_size=self.unlearn_bs, sampler=sampler)
|
||||
res_iter = iter(res_loader)
|
||||
|
||||
for _ in range(self.s1):
|
||||
h_estimate = [p.clone() for p in g]
|
||||
|
||||
for _ in range(self.s2):
|
||||
global_step += 1
|
||||
|
||||
try:
|
||||
data, target = next(res_iter)
|
||||
except StopIteration:
|
||||
res_iter = iter(res_loader)
|
||||
data, target = next(res_iter)
|
||||
|
||||
data, target = data.to(device), target.to(device)
|
||||
|
||||
# OPTIMIZATION 3: Clean up graph creation for loss & L2
|
||||
outputs = model(data)
|
||||
loss = criterion(outputs, target)
|
||||
|
||||
l2_reg_term = sum(p.pow(2).sum() for p in params)
|
||||
loss += (self.l2_reg + self.gamma) * l2_reg_term
|
||||
|
||||
h_s = self._hvp(loss, params, h_estimate)
|
||||
|
||||
# OPTIMIZATION 4: Avoid deprecated .data, use detach() and in-place ops
|
||||
with torch.no_grad():
|
||||
for k in range(len(params)):
|
||||
h_estimate[k].copy_(h_estimate[k] + g[k] - (h_s[k] / self.scale))
|
||||
|
||||
|
||||
# feed back on status
|
||||
if global_step % step_interval == 0 and current_pct < 100:
|
||||
current_pct += 1
|
||||
print(f"\rProgress: {current_pct}% done", end="", flush=True)
|
||||
|
||||
with torch.no_grad():
|
||||
for k in range(len(params)):
|
||||
h_res[k] += h_estimate[k] / self.scale
|
||||
|
||||
return [p / self.s1 for p in h_res]
|
||||
|
||||
|
||||
def _certify(self, model, device, delta, full_certification):
|
||||
certification = "full " if full_certification else "partial"
|
||||
print(f"Performing {certification} certification")
|
||||
|
||||
delta_idx = 0
|
||||
|
||||
# named_parameters to monitor layer positions
|
||||
for name, param in self.get_params(model, True):
|
||||
if param.requires_grad:
|
||||
noise = self.std * torch.randn(param.data.size(), device=device)
|
||||
|
||||
if full_certification:
|
||||
param.data.add_(delta[delta_idx] + noise)
|
||||
else:
|
||||
# option for applying certification only to last layers
|
||||
# deprecated
|
||||
if "layer4" in name or "fc" in name:
|
||||
param.data.add_(delta[delta_idx] + noise)
|
||||
else:
|
||||
# Keep early low-level vision filters entirely pristine
|
||||
pass
|
||||
|
||||
# Move to the next calculated Hessian vector block only after a valid update step
|
||||
delta_idx += 1
|
||||
|
||||
return model
|
||||
|
||||
|
||||
|
||||
def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
|
||||
device = next(model.parameters()).device
|
||||
|
||||
print(">> Calculating stable base gradients over the Forget set...")
|
||||
g = self._compute_loss_gradient(model, forget_loader, device)
|
||||
|
||||
print(">> Estimating non-convex inverse Hessian trajectories via Taylor series...")
|
||||
dataset = retain_loader.dataset
|
||||
delta = self._stochastic_newton_update(g, dataset, model, device)
|
||||
|
||||
print(">> Applying parameter removal adjustments (-delta)...")
|
||||
model = self._certify(
|
||||
model= model,
|
||||
device = device,
|
||||
delta = delta,
|
||||
full_certification = True
|
||||
)
|
||||
|
||||
print(">> Certified Unlearning process completed successfully.")
|
||||
return model
|
||||
|
||||
|
||||
# overriden function
|
||||
def _split_data(self, dataset):
|
||||
# Certified unlearning does require both forget and retain sets
|
||||
# split horizontaly. one class to forget and the rest to retain
|
||||
return get_unlearning_loaders(
|
||||
dataset=dataset,
|
||||
forget_class_idx=self.target_class_index,
|
||||
batch_size = 32
|
||||
)
|
||||
Reference in New Issue
Block a user