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, #num_classes = num_classes, 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 # 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() 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 )