188 lines
6.2 KiB
Python
188 lines
6.2 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()
|
|
# 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,
|
|
loader = forget_loader,
|
|
device = device
|
|
)
|
|
|
|
# 9
|
|
Z = self._compute_z(tensor=self.A, forget_index=forget_index)
|
|
|
|
# we transpose A for column based math
|
|
A_t = self.A.t()
|
|
# pi(a) drops the forget index logit dim( K x K-1)
|
|
a_pi = self.pi_mask(index=forget_index, tensor=self.A)
|
|
|
|
# construct B_z
|
|
B_r = a_pi.clone()
|
|
B_r[forget_index] = Z
|
|
B_z = B_r.t()
|
|
|
|
|
|
# A inverse
|
|
A_inv = torch.linalg.pinv(A_t, rcond=1e-2)
|
|
# 11
|
|
W_temp = B_z @ A_inv @ W
|
|
# Linear (Normalisation)filtration is done here,
|
|
# next steps are added for safe automated evaluations
|
|
|
|
# with one class removed, we have a head W_temp for 19 classes.
|
|
# but we have loaders for 20 classes for evaluation. so ,
|
|
# map K-1 W back to K x K.
|
|
W_z = torch.zeros_like(W)
|
|
mask = torch.ones(self.num_classes,dtype=torch.bool, device = device)
|
|
mask[forget_index] = False
|
|
W_z[mask] = W_temp
|
|
|
|
|
|
# 12
|
|
clf = self._get_classifier(model)
|
|
# load the weights
|
|
with torch.no_grad():
|
|
clf.weight.copy_(W_z)
|
|
# neutralise forget bias if exists
|
|
if clf.bias is not None:
|
|
n_bias = clf.bias.data.clone()
|
|
# push it down to negative
|
|
n_bias[forget_index] = -1e9
|
|
clf.bias.copy_(n_bias)
|
|
|
|
|
|
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
|
|
) |