136 lines
5.4 KiB
Python
136 lines
5.4 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,
|
|
num_classes: int = 20,
|
|
epochs: int = 6,
|
|
lr: float = 100.0,
|
|
gamma: float = 0.01,
|
|
):
|
|
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 = 25
|
|
|
|
|
|
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,
|
|
size=self.num_classes,
|
|
original_model=model,
|
|
target_class_index=self.target_class_index
|
|
)
|
|
|
|
# a WF_net module to be trained (unlearned) to generate alpha
|
|
wf_net = wf_model.get()
|
|
optimizer = optim.SGD([wf_net.alpha], lr=self.lr)
|
|
criterion = 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_rows = []
|
|
for label in r_labels:
|
|
allowed = [i for i in range(self.num_classes) if i != label.item()]
|
|
random_rows.append(np.random.choice(allowed))
|
|
|
|
gate_signals_r = torch.tensor(random_rows, dtype=torch.long, device=device)
|
|
outputs_r = wf_net(r_inputs, target_class_indices=gate_signals_r)
|
|
loss_r = criterion(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)
|
|
|
|
loss_f = 0.0
|
|
classes_in_batch = 0
|
|
|
|
# every image of class c will unlearn over the same row of alpha_l (poppi et al page 5)
|
|
for c in range(self.num_classes):
|
|
class_mask = (f_labels == c)
|
|
if not class_mask.any():
|
|
continue
|
|
|
|
labels_c = f_labels[class_mask]
|
|
|
|
# Slice the existing outputs instead of recalculating a forward pass
|
|
outputs_f_c = outputs_f[class_mask]
|
|
|
|
loss_f_ce = criterion(outputs_f_c, labels_c)
|
|
|
|
# Poppi et al. suggest employing reciprocal of the forget loss
|
|
# to avoid shortcomings of negative gradient approach
|
|
loss_f += 1.0 / (loss_f_ce + 1e-6)
|
|
classes_in_batch += 1
|
|
|
|
# Average forget loss by number of distinct classes seen in this batch
|
|
if classes_in_batch > 0:
|
|
loss_f = loss_f / classes_in_batch
|
|
|
|
# Regilarisation penality
|
|
loss_reg = torch.sum(1.0 - torch.sigmoid(wf_net.alpha))
|
|
|
|
# back propagation
|
|
total_loss = loss_r + (self.lambda_1 * loss_f) + (self.gamma * loss_reg)
|
|
total_loss.backward()
|
|
optimizer.step()
|
|
|
|
t_loss_r += loss_r.item()
|
|
t_loss_f += loss_f.item() if classes_in_batch > 0 else 0.0
|
|
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
|
|
)
|
|
|