Tinsae/Finetuning
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

certified

Browse Source
This commit is contained in:
Tinsae
2026-06-24 21:05:06 +02:00
parent 207fcae699
commit 3c6ee9e12d
10 changed files with 610 additions and 281 deletions
Download Patch File Download Diff File Expand all files Collapse all files

176
unlearning/WeightFiltration.py
View File

@@ -3,97 +3,34 @@ 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 WeightFiltration(Strategy):
"""
Implements Poppi et al.'s Weight Filtering framework for linear layers.
Uses a standard functional hook to guarantee native PyTorch autograd tracking.
Verbatim implementation of Poppi et al.'s WF-Net framework.
Directly filters the convolutional weights of a target layer using a learnable
channel mask, optimizing it via weight-space regularization.
"""
def __init__(self, target_class_index,num_classes: int, epochs: int = 10, lr: float = 0.2, gamma: float = 10.0):
super().__init__(target_class_index = target_class_index)
self.num_classes = num_classes
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
self.alpha = None
self.hook_handle = None
#self.alpha = None
def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
device = next(model.parameters()).device
model.eval()
def _optimise_filter(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader, device):
# 1. Initialize the wrapper with your pre-trained model
num_classes = model.fc.out_features
wf_model = WF_Net(original_model=model, num_classes=num_classes).to(device)
# Locate layer4 for dynamic optimization
target_layer = model.layer4 if hasattr(model, 'layer4') else None
fc_layer = model.fc if hasattr(model, 'fc') and isinstance(model.fc, nn.Linear) else None
if target_layer is None or fc_layer is None:
raise AttributeError("Model does not have the required layers.")
# Match alpha dimensions to the channels outputted by layer4
num_features = fc_layer.weight.shape[1]
self.alpha = nn.Parameter(torch.ones(self.num_classes, num_features, device=device) * 1.5)
# Freeze everything except our channel mask
for p in model.parameters():
p.requires_grad = False
self.alpha.requires_grad = True
# Hook into layer4 dynamically to run the untraining optimization
self.hook_handle = target_layer.register_forward_hook(self._get_hook())
# optimise the filter to maintain accuracy on retain set
# and decrease accuracy on forget set
self._optimise_filter(model, forget_loader, retain_loader, device)
# Remove the runtime hook
self.hook_handle.remove()
# Transfer the channel suppression permanently into model.fc
with torch.no_grad():
mask = torch.sigmoid(self.alpha[self.target_class_index]) # Shape: (num_features,)
# Suppress the channels ONLY for the target class row in fc
fc_layer.weight[self.target_class_index].copy_(
fc_layer.weight[self.target_class_index] * mask
)
print(f">> Baked deep channel filter into Class {self.target_class_index} weights.")
return model
def _get_hook(self):
"""
Filters the internal feature map channels of layer4.
The mask scales the channels across the batch.
"""
def functional_hook(module, layer_input, layer_output):
# layer_output shape: (batch, channels, height, width) -> e.g., (16, 2048, 7, 7)
# self.alpha shape: (num_classes, channels) -> e.g., (20, 2048)
# Extract 1D mask for the target class: (channels,)
mask = torch.sigmoid(self.alpha[self.target_class_index])
# Reshape mask to (1, channels, 1, 1) so it broadcasts over batch, height, and width
mask = mask.view(1, -1, 1, 1)
# Scale the internal feature maps before they move to the next layer
return layer_output * mask
return functional_hook
def _optimise_filter(self, model, forget_loader, retain_loader, device):
optimizer = optim.Adam([self.alpha], lr=self.lr)
# 2. ONLY optimize alpha (everything else is frozen inside the wrapper)
optimizer = optim.Adam([wf_model.alpha], lr=self.lr)
criterion = nn.CrossEntropyLoss()
print(f"[{self.__class__.__name__}] Unlearning Class {self.target_class_index} with gamma={self.gamma}...")
# To optimise this loop we will watch improvements after each optimisation
temp_forget_loss = None
# this can be adjusted to optimise the best escape point
# it is the value we set to evaluate performance improvement after each itteration.
# if improvement is less than this, then we break itteration.
threshold = 0.05
for epoch in range(self.epochs):
forget_iter = iter(forget_loader)
t_loss_r, t_loss_f = 0.0, 0.0
@@ -102,6 +39,7 @@ class WeightFiltration(Strategy):
for r_inputs, r_labels in retain_loader:
r_inputs, r_labels = r_inputs.to(device), r_labels.to(device)
# Pull the matching forget batch input
try:
f_inputs, _ = next(forget_iter)
except StopIteration:
@@ -111,10 +49,19 @@ class WeightFiltration(Strategy):
optimizer.zero_grad()
# Compute Losses
# The hook handles the weight filtering smoothly behind the scenes
loss_r = criterion(model(r_inputs), r_labels)
loss_f = -torch.sum((torch.ones_like(model(f_inputs)) / self.num_classes) * torch.log_softmax(model(f_inputs), dim=-1))
# --- APPLY ALGORITHM 1 FORWARD PASS TO BOTH INPUTS ---
# Pass the input batch AND the target unlearn class index
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)
# Compute Losses using Poppi et al.'s temperature scaled entropy
loss_r = criterion(outputs_r, r_labels)
temperature = 3.0
logits_f_scaled = outputs_f / temperature
loss_f = -torch.sum(
(torch.ones_like(logits_f_scaled) / num_classes) * torch.log_softmax(logits_f_scaled, dim=-1)
)
total_loss = loss_r + (self.gamma * loss_f)
total_loss.backward()
@@ -122,17 +69,56 @@ class WeightFiltration(Strategy):
t_loss_r += loss_r.item()
t_loss_f += loss_f.item()
steps += 1
forget_loss = t_loss_f / steps
print(f" Epoch {epoch+1}/{self.epochs} | Retain Loss: {t_loss_r/steps:.4f} | Forget Loss: {forget_loss:.4f}")
print(f" Epoch {epoch+1}/{self.epochs} | Retain Loss: {t_loss_r/steps:.4f} | Forget Loss: {t_loss_f/steps:.4f}")
if temp_forget_loss is not None:
return wf_model
improvement = temp_forget_loss - forget_loss
# if optimisation reaches a point of diminishing returns (improvements is less than threshold)
# we break the loop
if improvement < threshold:
break
# else we update the lasst recorded loss.
temp_forget_loss = forget_loss
def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
device = next(model.parameters()).device
model.eval()
# In WF-Net, the mask targets the last major convolutional block
# For ResNet-18, that is the final conv layer in layer4 block 1
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.")
# Store a pristine, non-grad copy of the original trained weights
# Shape of conv2.weight: (out_channels, in_channels, kernel_size, kernel_size) -> e.g., (512, 512, 3, 3)
original_weights = target_conv.weight.data.clone().detach()
out_channels = original_weights.shape[0]
# Initialize alpha gate vector matching Poppi et al.'s initialization range
# Shape: (out_channels,) -> acting directly as a filter-level gate
#self.alpha = nn.Parameter(torch.ones(out_channels, device=device) * 1.5)
# Freeze the global model graph; only optimize our filter parameter mask
for p in model.parameters():
p.requires_grad = False
#self.alpha.requires_grad = True
wf_model = self._optimise_filter(
model,
forget_loader=forget_loader,
retain_loader=retain_loader,
device=device,
)
# --- PERMANENT BAKING STEP ---
# Disconnect the dynamic parameter and freeze the optimal gated state permanently into the architecture
with torch.no_grad():
final_mask = torch.sigmoid(wf_model.alpha[self.target_class_index]).view(-1, 1, 1, 1)
target_conv.weight.copy_(original_weights * final_mask)
# Re-enable model parameters for downstream evaluation processing
for p in model.parameters():
p.requires_grad = True
print(f">> Permanently altered {out_channels} convolutional filters in layer4 via WF-Net.")
return model