177 lines
5.9 KiB
Python
177 lines
5.9 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from .Strategy import Strategy
|
|
from torch.utils.data import DataLoader
|
|
from sets.Data import get_unlearning_loaders, _combine_set, vertical_split
|
|
|
|
class LinearFiltration(Strategy):
|
|
def __init__(self, target_class_index, num_classes = 20):
|
|
super().__init__(target_class_index=target_class_index)
|
|
self.A = None
|
|
self.num_classes = num_classes
|
|
|
|
def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
|
|
model.eval()
|
|
device = next(model.parameters()).device
|
|
|
|
return self.normalise(
|
|
model=model,
|
|
retain_loader=retain_loader,
|
|
forget_loader=forget_loader,
|
|
device=device,
|
|
forget_index=self.target_class_index
|
|
)
|
|
|
|
|
|
def _get_classifier(self, model: nn.Module) -> nn.Linear:
|
|
|
|
inner_model = getattr(model, "model", model)
|
|
|
|
# looking for standard naming conventions in named modules
|
|
for name, module in inner_model.named_modules():
|
|
# Check if it's our target linear layer
|
|
if (name == "fc" or name == "classifier") and isinstance(module, nn.Linear):
|
|
return module
|
|
|
|
# Handle models (like EfficientNet) where the classifier is a Sequential block
|
|
if name == "classifier" and isinstance(module, nn.Sequential):
|
|
for sub_module in reversed(list(module.children())):
|
|
if isinstance(sub_module, nn.Linear):
|
|
return sub_module
|
|
|
|
# scan backwards for the last Linear layer
|
|
for module in reversed(list(inner_model.modules())):
|
|
if isinstance(module, nn.Linear):
|
|
return module
|
|
|
|
raise RuntimeError(f"Could not locate a linear classification head for {model.__class__.__name__}")
|
|
|
|
|
|
|
|
def _compute_A(self, model, loader, device):
|
|
model.eval()
|
|
|
|
|
|
# Initialize tracking tensors
|
|
sums = torch.zeros(self.num_classes, self.num_classes, device=device)
|
|
counts = torch.zeros(self.num_classes, device=device)
|
|
|
|
with torch.no_grad():
|
|
for inputs, targets in loader:
|
|
inputs, targets = inputs.to(device), targets.to(device)
|
|
|
|
# the logit predictions
|
|
outputs = model(inputs)
|
|
|
|
# One-hot encode targets to act as a routing mask
|
|
one_hot = torch.nn.functional.one_hot(targets, num_classes=self.num_classes).float()
|
|
|
|
# add
|
|
sums += torch.t(one_hot) @ outputs
|
|
|
|
# Sum columns of one-hot to get counts per class in this batch
|
|
counts += one_hot.sum(dim=0)
|
|
|
|
# means
|
|
counts_safe = counts.unsqueeze(1)
|
|
self.A = torch.where(
|
|
counts_safe > 0,
|
|
sums / counts_safe,
|
|
torch.zeros_like(sums)
|
|
)
|
|
|
|
def pi_mask(self, index, tensor):
|
|
mask = torch.ones(self.num_classes, dtype= torch.bool,device = tensor.device)
|
|
mask[index] = False
|
|
return tensor[:, mask]
|
|
|
|
# 9
|
|
def _compute_z(self, tensor, forget_index):
|
|
|
|
K = tensor.shape[0]
|
|
a_pi = self.pi_mask(tensor = tensor, index = forget_index)
|
|
|
|
|
|
#pi_a_f = torch.zeros(tensor.shape[1], device=tensor.device)
|
|
|
|
t_1 = a_pi[forget_index] #pi_a_f
|
|
# row vector for the forgotten class
|
|
#a_f = tensor[forget_index, :]
|
|
|
|
|
|
# We compute the target shift over features
|
|
t_2 = (1.0 / (K - 1)) * t_1.sum()
|
|
|
|
mask = torch.ones(self.num_classes, dtype= torch.bool,device = tensor.device)
|
|
mask[forget_index] = False
|
|
remaining_rows = a_pi[mask]
|
|
|
|
#r_A = tensor[mask_rows, :]
|
|
t_3 = (1.0 / ((K - 1)) ** 2) * remaining_rows.sum()
|
|
|
|
return t_1 - t_2 + t_3
|
|
|
|
|
|
|
|
# Normalisation filtration
|
|
def normalise(self, model, retain_loader, forget_loader, device, forget_index):
|
|
clf = self._get_classifier(model)
|
|
W = clf.weight.data.clone()
|
|
#num_classes = W.shape[0]
|
|
|
|
# we combine the data so we can calculate the mean of prdictions
|
|
#full_loader = _combine_set(retain_loader, forget_loader)
|
|
# 8
|
|
# Computing A is the most resource intensive part of this algorithm
|
|
# and to optimise the process, we computr it only once and re-use it
|
|
# because mean of all prdictions is the same for all
|
|
if self.A is None:
|
|
self._compute_A(
|
|
model = model,
|
|
#num_classes = num_classes,
|
|
loader = forget_loader,
|
|
device = device
|
|
)
|
|
|
|
# 9
|
|
Z = self._compute_z(tensor=self.A, forget_index=forget_index)
|
|
B_Z_rows = []
|
|
|
|
a_pi = self.pi_mask(index=forget_index, tensor=self.A)
|
|
|
|
for i in range(self.num_classes):
|
|
if i == forget_index:
|
|
B_Z_rows.append(Z)
|
|
else:
|
|
# Retained classes maintain their original ideal feature directions
|
|
B_Z_rows.append(a_pi[i])
|
|
|
|
# 10
|
|
# Stack back along dim=0 to match (num_classes, h_dim)
|
|
# to get mean
|
|
B_Z = torch.stack(B_Z_rows, dim=0)
|
|
|
|
A_inv = torch.linalg.pinv(a_pi)
|
|
# 11
|
|
W_Z = B_Z @ A_inv @ W
|
|
|
|
# 12
|
|
clf = self._get_classifier(model)
|
|
with torch.no_grad():
|
|
clf.weight.copy_(W_Z)
|
|
|
|
return model
|
|
|
|
# overriden function
|
|
def _split_data(self, dataset):
|
|
'''return get_unlearning_loaders(
|
|
dataset=dataset,
|
|
forget_class_idx=self.target_class_index,
|
|
batch_size = 32
|
|
)'''
|
|
return vertical_split(
|
|
dataset= dataset,
|
|
batch_size=32,
|
|
num_classes=self.num_classes,
|
|
ratio=0.1
|
|
) |