213 lines
8.0 KiB
Python
213 lines
8.0 KiB
Python
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)
|
|
|
|
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]
|
|
|
|
# Data mean conversion
|
|
for i in range(len(grad_accumulator)):
|
|
grad_accumulator[i] /= total_samples
|
|
|
|
# regularisation gradient
|
|
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]
|
|
# hesian estimation
|
|
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)
|
|
|
|
# forward
|
|
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
|
|
) |