155 lines
4.4 KiB
Python
155 lines
4.4 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
|
|
|
|
class LinearFiltration(Strategy):
|
|
def __init__(self, target_class_index):
|
|
super().__init__(target_class_index=target_class_index)
|
|
|
|
def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
|
|
model.eval()
|
|
# Freeze internal params
|
|
for param in model.parameters():
|
|
param.requires_grad = False
|
|
|
|
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 _sums_and_counts(self, model, num_classes, loader, device, forget_index, h_dim):
|
|
model.eval()
|
|
|
|
sums = torch.zeros(num_classes, h_dim, device=device)
|
|
counts = torch.zeros(num_classes, device=device)
|
|
|
|
# Generate values for retain
|
|
with torch.no_grad():
|
|
for inputs, targets in loader:
|
|
inputs = inputs.to(device)
|
|
targets = targets.to(device)
|
|
# predictions
|
|
outputs = model(inputs)
|
|
|
|
for j in range(num_classes):
|
|
if j == forget_index:
|
|
continue
|
|
mask = (targets == j)
|
|
|
|
if mask.any():
|
|
sums[j] += outputs[mask].sum(dim=0)
|
|
counts[j] += mask.sum()
|
|
|
|
return sums, counts
|
|
|
|
|
|
#
|
|
def _get_means(self,model, num_classes, loader, device, forget_index):
|
|
h_dim = model.fc.out_features
|
|
|
|
# all predictions
|
|
sums, counts = self._sums_and_counts(
|
|
model=model,
|
|
num_classes=num_classes,
|
|
loader=loader,
|
|
device=device,
|
|
forget_index=forget_index,
|
|
h_dim=h_dim
|
|
)
|
|
|
|
#A = []
|
|
|
|
counts_safe = counts.unsqueeze(1)
|
|
A = torch.where(
|
|
counts_safe > 0,
|
|
sums / counts_safe,
|
|
torch.zeros_like(sums)
|
|
)
|
|
# 6
|
|
return A
|
|
|
|
|
|
# 9
|
|
def _compute_z(self, tensor, forget_index):
|
|
|
|
K = tensor.shape[0]
|
|
|
|
# pi_a_forget should match the feature space dimensions (h_dim)
|
|
pi_a_f = torch.zeros(tensor.shape[1], device=tensor.device)
|
|
|
|
t_1 = pi_a_f
|
|
# Extracting the row vector for the forgotten class
|
|
a_f = tensor[forget_index, :]
|
|
|
|
mask_a_f = torch.ones(
|
|
a_f.shape[0],
|
|
dtype=torch.bool,
|
|
device=tensor.device
|
|
)
|
|
# We compute the target shift over features
|
|
t_2 = -(1.0 / (K - 1)) * a_f[mask_a_f].sum()
|
|
|
|
mask_rows = torch.ones(K, dtype=torch.bool, device=tensor.device)
|
|
mask_rows[forget_index] = False
|
|
|
|
r_A = tensor[mask_rows, :]
|
|
t_3 = (1.0 / ((K - 1)) ** 2) * r_A.sum()
|
|
|
|
return t_1 + t_2 + t_3
|
|
|
|
|
|
# Normalisation filtration
|
|
def normalise(self, model, retain_loader, forget_loader, device, forget_index):
|
|
W = model.fc.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
|
|
A = self._get_means(
|
|
model=model,
|
|
num_classes=num_classes,
|
|
loader=full_loader,
|
|
device=device,
|
|
forget_index=forget_index
|
|
)
|
|
# 9
|
|
Z = self._compute_z(tensor=A, forget_index=forget_index)
|
|
B_Z_rows = []
|
|
|
|
for i in range(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[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)
|
|
# 11
|
|
W_Z = B_Z @ A_inv @ W
|
|
|
|
# 12
|
|
model.fc.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
|
|
) |