138 lines
5.6 KiB
Python
138 lines
5.6 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.optim as optim
|
|
from torch.utils.data import DataLoader, ConcatDataset, Subset
|
|
from unlearning.Strategy import Strategy
|
|
import numpy as np
|
|
from sklearn.metrics import classification_report
|
|
from architectures.WFNet import WF_Net_Model
|
|
|
|
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
|
|
):
|
|
super().__init__(target_class_index=target_class_index)
|
|
self.epochs = epochs
|
|
self.lr = lr
|
|
self.gamma = gamma
|
|
self.num_classes = num_classes
|
|
self.wf_model = None
|
|
self.lambda_1 = lambda_1
|
|
self.arch = arch
|
|
|
|
|
|
def _optimise_filter(self, model: nn.Module, retain_loader: DataLoader, forget_loader: DataLoader, device) -> nn.Module:
|
|
|
|
# new WF_Model instance
|
|
wf_model = WF_Net_Model(
|
|
device=device,
|
|
arch=self.arch,
|
|
size=self.num_classes,
|
|
original_model=model,
|
|
target_class_index=self.target_class_index
|
|
)
|
|
|
|
wf_net = wf_model.get()
|
|
optimizer = optim.SGD([wf_net.alpha], lr=self.lr)
|
|
# Use reduction='none' so we can manipulate individual item losses
|
|
criterion_none = nn.CrossEntropyLoss(reduction='none')
|
|
criterion_mean = nn.CrossEntropyLoss()
|
|
|
|
for epoch in range(self.epochs):
|
|
t_loss_r, t_loss_f = 0.0, 0.0
|
|
steps = 0
|
|
# forget and retain
|
|
for (r_inputs, r_labels), (f_inputs, f_labels) in zip(retain_loader, forget_loader):
|
|
r_inputs, r_labels = r_inputs.to(device), r_labels.to(device)
|
|
f_inputs, f_labels = f_inputs.to(device), f_labels.to(device)
|
|
|
|
optimizer.zero_grad()
|
|
|
|
# retain data paired with randomly selected rows of alpha to compute the retaining loss
|
|
random_offset = torch.randint(0, self.num_classes - 1, size=r_labels.shape, device=device)
|
|
gate_signals_r = torch.where(random_offset >= r_labels, random_offset + 1, random_offset)
|
|
|
|
outputs_r = wf_net(r_inputs, target_class_indices=gate_signals_r)
|
|
loss_r = criterion_mean(outputs_r, r_labels)
|
|
|
|
# Forget set is paired with corresponding labels as row selectors for alpha
|
|
# and used to compute unlearning loss
|
|
outputs_f = wf_net(f_inputs, target_class_indices=f_labels)
|
|
|
|
# Calculate loss for every single item in the batch at once
|
|
per_item_forget_loss = criterion_none(outputs_f, f_labels)
|
|
|
|
# Use a scatter/sum approach to get class-wise losses without a Python loop
|
|
# Create a mask of unique classes present in this batch
|
|
unique_classes, inverse_indices = torch.unique(f_labels, return_inverse=True)
|
|
classes_in_batch = unique_classes.size(0)
|
|
|
|
if classes_in_batch > 0:
|
|
# average CE loss per class
|
|
class_loss_sums = torch.zeros(classes_in_batch, device=device)
|
|
class_loss_sums.scatter_add_(0, inverse_indices, per_item_forget_loss)
|
|
|
|
class_counts = torch.zeros(classes_in_batch, device=device)
|
|
class_counts.scatter_add_(0, inverse_indices, torch.ones_like(per_item_forget_loss))
|
|
|
|
mean_class_ce_loss = class_loss_sums / class_counts
|
|
|
|
# Poppi et al. suggest employing reciprocal of the forget loss
|
|
# to avoid shortcomings of negative gradient approach
|
|
loss_f = torch.mean(1.0 / (mean_class_ce_loss + 1e-6))
|
|
else:
|
|
loss_f = torch.tensor(0.0, device=device)
|
|
|
|
# Regularisation penalty
|
|
loss_reg = torch.sum(1.0 - torch.sigmoid(wf_net.alpha))
|
|
|
|
# Backpropagation
|
|
total_loss = loss_r + (self.lambda_1 * loss_f) + (self.gamma * loss_reg)
|
|
total_loss.backward()
|
|
optimizer.step()
|
|
|
|
# Keep tracking stats
|
|
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()
|
|
|
|
if self.wf_model is None:
|
|
print(">> Initializing and compiling global WF-Net matrix (Run Once for all classes)...")
|
|
|
|
self.wf_model = self._optimise_filter(
|
|
model,
|
|
retain_loader=retain_loader,
|
|
forget_loader=forget_loader,
|
|
device=device
|
|
)
|
|
else:
|
|
print(f">> Gating matrix loaded. Switching layout to target class index: {self.target_class_index}")
|
|
self.wf_model.target_class_index = self.target_class_index
|
|
|
|
return self.wf_model
|
|
|
|
def _split_data(self, dataset):
|
|
return vertical_split(
|
|
dataset= dataset,
|
|
batch_size=32,
|
|
num_classes=self.num_classes
|
|
)
|
|
|