127 lines
5.5 KiB
Python
127 lines
5.5 KiB
Python
import torch
|
|
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):
|
|
"""
|
|
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: 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
|
|
|
|
|
|
|
|
|
|
|
|
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)
|
|
|
|
# 2. ONLY optimize alpha (everything else is frozen inside the wrapper)
|
|
optimizer = optim.Adam([wf_model.alpha], lr=self.lr)
|
|
criterion = nn.CrossEntropyLoss()
|
|
|
|
for epoch in range(self.epochs):
|
|
forget_iter = iter(forget_loader)
|
|
t_loss_r, t_loss_f = 0.0, 0.0
|
|
steps = 0
|
|
|
|
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:
|
|
forget_iter = iter(forget_loader)
|
|
f_inputs, _ = next(forget_iter)
|
|
f_inputs = f_inputs.to(device)
|
|
|
|
optimizer.zero_grad()
|
|
|
|
# --- 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
|
|
|
|
# Compute uniform target entropy per-sample, then average over the batch
|
|
log_probs_f = torch.log_softmax(logits_f_scaled, dim=-1)
|
|
uniform_target = torch.ones_like(logits_f_scaled) / num_classes
|
|
loss_f = -torch.sum(uniform_target * log_probs_f, dim=-1).mean()
|
|
|
|
total_loss = loss_r + (self.gamma * loss_f)
|
|
total_loss.backward()
|
|
optimizer.step()
|
|
|
|
t_loss_r += loss_r.item()
|
|
t_loss_f += loss_f.item()
|
|
steps += 1
|
|
|
|
print(f" Epoch {epoch+1}/{self.epochs} | Retain Loss: {t_loss_r/steps:.4f} | Forget Loss: {t_loss_f/steps:.4f}")
|
|
return wf_model
|
|
|
|
|
|
|
|
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 ---
|
|
with torch.no_grad():
|
|
# Grab the alpha mask vector for the forgotten class and cast to 4D tensor shape
|
|
final_mask = torch.sigmoid(wf_model.alpha[self.target_class_index]).view(-1, 1, 1, 1)
|
|
|
|
# Apply filter masking permanently back onto the base layer
|
|
target_conv.weight.copy_(original_weights * final_mask)
|
|
|
|
# Unfreeze architecture parameters for evaluations downstream
|
|
for p in model.parameters():
|
|
p.requires_grad = True
|
|
|
|
print(f">> Permanently altered {out_channels} convolutional filters in layer4 via WF-Net.")
|
|
return model
|