added reports and params
This commit is contained in:
@@ -14,9 +14,18 @@ class CertifiedUnlearning(Strategy):
|
||||
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):
|
||||
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
|
||||
@@ -49,151 +58,12 @@ class CertifiedUnlearning(Strategy):
|
||||
|
||||
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
|
||||
|
||||
@@ -208,12 +78,11 @@ class CertifiedUnlearning(Strategy):
|
||||
for i in range(len(grad_accumulator)):
|
||||
grad_accumulator[i] += mini_grads[i]
|
||||
|
||||
# Empirical data mean conversion
|
||||
# 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
|
||||
# regularisation gradient
|
||||
for i, param in enumerate(params):
|
||||
grad_accumulator[i] += 2 * self.l2_reg * param.detach()
|
||||
|
||||
@@ -243,7 +112,7 @@ class CertifiedUnlearning(Strategy):
|
||||
|
||||
for _ in range(self.s1):
|
||||
h_estimate = [p.clone() for p in g]
|
||||
|
||||
# hesian estimation
|
||||
for _ in range(self.s2):
|
||||
global_step += 1
|
||||
|
||||
@@ -255,7 +124,7 @@ class CertifiedUnlearning(Strategy):
|
||||
|
||||
data, target = data.to(device), target.to(device)
|
||||
|
||||
# OPTIMIZATION 3: Clean up graph creation for loss & L2
|
||||
# forward
|
||||
outputs = model(data)
|
||||
loss = criterion(outputs, target)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user