unlearning done
This commit is contained in:
@@ -2,6 +2,7 @@ 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):
|
||||
@@ -23,40 +24,8 @@ class LinearFiltration(Strategy):
|
||||
forget_index=self.target_class_index
|
||||
)
|
||||
|
||||
# FIX: Added staticmethod decorator
|
||||
@staticmethod
|
||||
def get_features(model, inputs):
|
||||
# For ResNet, pass through everything up to the fc layer
|
||||
x = model.conv1(inputs)
|
||||
x = model.bn1(x)
|
||||
x = model.relu(x)
|
||||
x = model.maxpool(x)
|
||||
|
||||
x = model.layer1(x)
|
||||
x = model.layer2(x)
|
||||
x = model.layer3(x)
|
||||
x = model.layer4(x)
|
||||
|
||||
x = model.avgpool(x)
|
||||
x = torch.flatten(x, 1)
|
||||
return x
|
||||
|
||||
@staticmethod
|
||||
def _calculate_filtration_matrix(num_classes: int, forget_class: int, device: torch.device) -> torch.Tensor:
|
||||
A = torch.eye(num_classes, device=device)
|
||||
num_remaining = num_classes - 1
|
||||
|
||||
for j in range(num_classes):
|
||||
if j == forget_class:
|
||||
A[forget_class, j] = 0.0
|
||||
else:
|
||||
A[forget_class, j] = 1.0 / num_remaining
|
||||
|
||||
return A
|
||||
|
||||
|
||||
@staticmethod
|
||||
def _sums_and_counts(model, num_classes, retain_loader, forget_loader, device, forget_index, h_dim):
|
||||
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)
|
||||
@@ -64,11 +33,11 @@ class LinearFiltration(Strategy):
|
||||
|
||||
# Generate values for retain
|
||||
with torch.no_grad():
|
||||
for inputs, targets in retain_loader:
|
||||
for inputs, targets in loader:
|
||||
inputs = inputs.to(device)
|
||||
targets = targets.to(device)
|
||||
# FIX: Call get_features instead of model() directly
|
||||
outputs = LinearFiltration.get_features(model, inputs)
|
||||
# predictions
|
||||
outputs = model(inputs)
|
||||
|
||||
for j in range(num_classes):
|
||||
if j == forget_index:
|
||||
@@ -79,65 +48,54 @@ class LinearFiltration(Strategy):
|
||||
sums[j] += outputs[mask].sum(dim=0)
|
||||
counts[j] += mask.sum()
|
||||
|
||||
# Values for forget
|
||||
with torch.no_grad():
|
||||
for inputs, targets in forget_loader:
|
||||
inputs = inputs.to(device)
|
||||
targets = targets.to(device)
|
||||
# FIX: Call get_features instead of model() directly
|
||||
outputs = LinearFiltration.get_features(model, inputs)
|
||||
|
||||
mask = (targets == forget_index)
|
||||
|
||||
if mask.any():
|
||||
sums[forget_index] += outputs[mask].sum(dim=0)
|
||||
counts[forget_index] += mask.sum()
|
||||
|
||||
return sums, counts
|
||||
|
||||
@staticmethod
|
||||
def _get_means(model, num_classes, retain_loader, forget_loader, device, forget_index):
|
||||
h_dim = model.fc.in_features
|
||||
|
||||
sums, counts = LinearFiltration._sums_and_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,
|
||||
retain_loader=retain_loader,
|
||||
forget_loader=forget_loader,
|
||||
loader=loader,
|
||||
device=device,
|
||||
forget_index=forget_index,
|
||||
h_dim=h_dim
|
||||
)
|
||||
A = []
|
||||
|
||||
for i in range(num_classes):
|
||||
if counts[i] > 0:
|
||||
A.append(sums[i] / counts[i])
|
||||
else:
|
||||
A.append(torch.zeros(h_dim, device=device))
|
||||
|
||||
# CORRECT: Stack along dim=0 to make it (num_classes, h_dim)
|
||||
return torch.stack(A, dim=0)
|
||||
#A = []
|
||||
|
||||
counts_safe = counts.unsqueeze(1)
|
||||
A = torch.where(
|
||||
counts_safe > 0,
|
||||
sums / counts_safe,
|
||||
torch.zeros_like(sums)
|
||||
)
|
||||
# 6
|
||||
return A
|
||||
|
||||
|
||||
@staticmethod
|
||||
def _compute_z(tensor, forget_index):
|
||||
# Now tensor has shape (num_classes, h_dim) -> tensor.shape[0] is num_classes
|
||||
# 9
|
||||
def _compute_z(self, tensor, forget_index):
|
||||
|
||||
K = tensor.shape[0]
|
||||
|
||||
# pi_a0 should match the feature space dimensions (h_dim)
|
||||
pi_a0 = torch.zeros(tensor.shape[1], device=tensor.device)
|
||||
# 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_a0
|
||||
a0 = tensor[forget_index, :] # Extracting the row vector for the forgotten class
|
||||
t_1 = pi_a_f
|
||||
# Extracting the row vector for the forgotten class
|
||||
a_f = tensor[forget_index, :]
|
||||
|
||||
mask_a0 = torch.ones(
|
||||
a0.shape[0],
|
||||
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)) * a0[mask_a0].sum()
|
||||
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
|
||||
@@ -148,21 +106,23 @@ class LinearFiltration(Strategy):
|
||||
return t_1 + t_2 + t_3
|
||||
|
||||
|
||||
@staticmethod
|
||||
def normalise(model, retain_loader, forget_loader, device, forget_index):
|
||||
# Normalisation filtration
|
||||
def normalise(self, model, retain_loader, forget_loader, device, forget_index):
|
||||
W = model.fc.weight.data.clone()
|
||||
num_classes = W.shape[0]
|
||||
|
||||
A = LinearFiltration._get_means(
|
||||
|
||||
# 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,
|
||||
retain_loader=retain_loader,
|
||||
forget_loader=forget_loader,
|
||||
loader=full_loader,
|
||||
device=device,
|
||||
forget_index=forget_index
|
||||
)
|
||||
|
||||
Z = LinearFiltration._compute_z(tensor=A, forget_index=forget_index)
|
||||
# 9
|
||||
Z = self._compute_z(tensor=A, forget_index=forget_index)
|
||||
B_Z_rows = []
|
||||
|
||||
for i in range(num_classes):
|
||||
@@ -172,13 +132,24 @@ class LinearFiltration(Strategy):
|
||||
# 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
|
||||
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
|
||||
)
|
||||
Reference in New Issue
Block a user