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)
|
||||
|
||||
|
||||
79
unlearning/Retrain.py
Normal file
79
unlearning/Retrain.py
Normal file
@@ -0,0 +1,79 @@
|
||||
import time
|
||||
from pathlib import Path
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.optim as optim
|
||||
from torch.utils.data import DataLoader
|
||||
from unlearning.Strategy import Strategy
|
||||
|
||||
class Retrain(Strategy):
|
||||
"""
|
||||
Implements the Exact Unlearning Baseline by retraining the model architecture
|
||||
completely from scratch using only the retained dataset partition.
|
||||
"""
|
||||
def __init__(self, target_class_index: int, lr: float = 0.01,
|
||||
weight_decay: float = 0.0005, epochs: int = 5):
|
||||
super().__init__(target_class_index)
|
||||
self.lr = lr
|
||||
self.weight_decay = weight_decay
|
||||
self.epochs = epochs
|
||||
|
||||
def _reset_weights(self, model: nn.Module):
|
||||
"""
|
||||
Re-initializes all learnable parameters of the model to clear pre-trained memories.
|
||||
"""
|
||||
inner_model = getattr(model, "model", model)
|
||||
for layer in inner_model.modules():
|
||||
if hasattr(layer, 'reset_parameters'):
|
||||
layer.reset_parameters()
|
||||
|
||||
def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
|
||||
device = next(model.parameters()).device
|
||||
|
||||
print(">> Triggering Exact Unlearning Baseline (Retraining from scratch)...")
|
||||
|
||||
# 1. Clear the pre-trained state completely
|
||||
self._reset_weights(model) # model should be loaded here or weights reset to ImageNet (pretrained default)
|
||||
model.train()
|
||||
|
||||
# fresh optimizer for this clean environment
|
||||
optimizer = optim.SGD(
|
||||
model.parameters(),
|
||||
lr=self.lr,
|
||||
momentum=0.9,
|
||||
weight_decay=self.weight_decay
|
||||
)
|
||||
criterion = nn.CrossEntropyLoss()
|
||||
|
||||
# 3. Standard training loop over the Retain set
|
||||
for epoch in range(self.epochs):
|
||||
running_loss = 0.0
|
||||
total_samples = 0
|
||||
|
||||
for data, targets in retain_loader:
|
||||
data, targets = data.to(device), targets.to(device)
|
||||
|
||||
optimizer.zero_grad()
|
||||
outputs = model(data)
|
||||
loss = criterion(outputs, targets)
|
||||
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
running_loss += loss.item() * targets.size(0)
|
||||
total_samples += targets.size(0)
|
||||
|
||||
epoch_loss = running_loss / total_samples
|
||||
print(f" [Retrain] Epoch {epoch+1}/{self.epochs} completed. Loss: {epoch_loss:.4f}")
|
||||
|
||||
print(">> Retraining pipeline finished. Baseline weights established.")
|
||||
return model
|
||||
|
||||
def _split_data(self, dataset):
|
||||
# Dynamically pulls loaders from your Data.py script
|
||||
from sets.Data import get_unlearning_loaders
|
||||
return get_unlearning_loaders(
|
||||
dataset=dataset,
|
||||
forget_class_idx=self.target_class_index,
|
||||
batch_size=32
|
||||
)
|
||||
@@ -11,15 +11,14 @@ from sets.Data import vertical_split
|
||||
|
||||
class WeightFiltration(Strategy):
|
||||
def __init__(self,
|
||||
target_class_index: int,
|
||||
arch,
|
||||
num_classes: int = 20,
|
||||
epochs: int = 6,
|
||||
lr: float = 100.0,
|
||||
gamma: float = 0.01,
|
||||
lambda_1 = 25
|
||||
|
||||
):
|
||||
target_class_index: int,
|
||||
arch,
|
||||
num_classes: int = 20,
|
||||
epochs: int = 6,
|
||||
lr: float = 100.0,
|
||||
gamma: float = 0.01,
|
||||
lambda_1 = 25
|
||||
):
|
||||
super().__init__(target_class_index=target_class_index)
|
||||
self.epochs = epochs
|
||||
self.lr = lr
|
||||
|
||||
@@ -1,86 +0,0 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
class WF_Net(nn.Module):
|
||||
"""
|
||||
Implements Poppi et al.'s WF Model structure.
|
||||
Wraps a pre-trained ResNet-18 and dynamically applies
|
||||
weight-space gating matrix multiplication during the forward step.
|
||||
"""
|
||||
def __init__(self, original_model: nn.Module, num_classes: int):
|
||||
super().__init__()
|
||||
# Extract the sequence of blocks/layers L from the original model
|
||||
self.conv1 = original_model.conv1
|
||||
self.bn1 = original_model.bn1
|
||||
self.relu = original_model.relu
|
||||
self.maxpool = original_model.maxpool
|
||||
|
||||
self.layer1 = original_model.layer1
|
||||
self.layer2 = original_model.layer2
|
||||
self.layer3 = original_model.layer3
|
||||
self.layer4 = original_model.layer4
|
||||
|
||||
self.avgpool = original_model.avgpool
|
||||
self.fc = original_model.fc
|
||||
|
||||
# Target layer for filtering: layer4 block 1 conv2
|
||||
# We extract its static tensor data out of the autograd parameter pool
|
||||
self.target_conv = self.layer4[1].conv2
|
||||
self.original_w = nn.Parameter(self.target_conv.weight.data.clone().detach(), requires_grad=False)
|
||||
|
||||
# Require: Alpha gating matrix. Shape: (num_classes, out_channels)
|
||||
# Initialized to 1.5 as per Poppi et al.'s verbatim specification
|
||||
out_channels = self.original_w.shape[0]
|
||||
#self.alpha = nn.Parameter(torch.ones(num_classes, out_channels) * 1.5)
|
||||
self.alpha = nn.Parameter(torch.ones(num_classes, out_channels))
|
||||
|
||||
def forward(self, x: torch.Tensor, target_class_indices: torch.Tensor) -> torch.Tensor:
|
||||
# 1. Run through early sequence of layers undisturbed
|
||||
x = self.maxpool(self.relu(self.bn1(self.conv1(x))))
|
||||
x = self.layer1(x)
|
||||
x = self.layer2(x)
|
||||
x = self.layer3(x)
|
||||
|
||||
# Run layer4 block 0 normally
|
||||
x = self.layer4[0](x)
|
||||
|
||||
# -------------------------------------------------------------
|
||||
# HERE IT IS: Save the structural skip connection (identity)
|
||||
# BEFORE modifying features via block 1's convolutions
|
||||
# -------------------------------------------------------------
|
||||
identity = x
|
||||
|
||||
# Now enter layer4 block 1
|
||||
x = self.layer4[1].conv1(x)
|
||||
x = self.layer4[1].bn1(x)
|
||||
x = self.layer4[1].relu(x)
|
||||
|
||||
# [Your Step 1 Masking Math happens right here...]
|
||||
batch_alpha = self.alpha[target_class_indices]
|
||||
mask = torch.sigmoid(batch_alpha).view(x.size(0), -1, 1, 1)
|
||||
|
||||
# Run the functional convolution
|
||||
x = F.conv2d(
|
||||
x,
|
||||
weight=self.original_w,
|
||||
bias=self.target_conv.bias,
|
||||
stride=self.target_conv.stride,
|
||||
padding=self.target_conv.padding
|
||||
)
|
||||
|
||||
# Apply your WF-Net channel mask
|
||||
x = x * mask
|
||||
x = self.layer4[1].bn2(x)
|
||||
|
||||
# -------------------------------------------------------------
|
||||
# HERE IT IS USED: Add the pristine identity back to the gated output
|
||||
# -------------------------------------------------------------
|
||||
x = self.layer4[1].relu(x + identity)
|
||||
|
||||
# Final Classification Head Sequence
|
||||
x = self.avgpool(x)
|
||||
x = torch.flatten(x, 1)
|
||||
y_out = self.fc(x)
|
||||
|
||||
return y_out
|
||||
Reference in New Issue
Block a user