reports from optimised linear filtration
This commit is contained in:
@@ -2,12 +2,13 @@ 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
|
||||
from sets.Data import get_unlearning_loaders, _combine_set, vertical_split
|
||||
|
||||
class LinearFiltration(Strategy):
|
||||
def __init__(self, target_class_index):
|
||||
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()
|
||||
@@ -51,13 +52,13 @@ class LinearFiltration(Strategy):
|
||||
|
||||
|
||||
|
||||
def _compute_A(self, model, num_classes, loader, device):
|
||||
def _compute_A(self, model, loader, device):
|
||||
model.eval()
|
||||
|
||||
|
||||
# Initialize tracking tensors
|
||||
sums = torch.zeros(num_classes, num_classes, device=device)
|
||||
counts = torch.zeros(num_classes, device=device)
|
||||
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:
|
||||
@@ -67,7 +68,7 @@ class LinearFiltration(Strategy):
|
||||
outputs = model(inputs)
|
||||
|
||||
# One-hot encode targets to act as a routing mask
|
||||
one_hot = torch.nn.functional.one_hot(targets, num_classes=num_classes).float()
|
||||
one_hot = torch.nn.functional.one_hot(targets, num_classes=self.num_classes).float()
|
||||
|
||||
# add
|
||||
sums += torch.t(one_hot) @ outputs
|
||||
@@ -77,7 +78,6 @@ class LinearFiltration(Strategy):
|
||||
|
||||
# means
|
||||
counts_safe = counts.unsqueeze(1)
|
||||
print(f"COUNTS IS >>>>>>>>> {counts_safe}")
|
||||
self.A = torch.where(
|
||||
counts_safe > 0,
|
||||
sums / counts_safe,
|
||||
@@ -117,10 +117,10 @@ class LinearFiltration(Strategy):
|
||||
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]
|
||||
#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)
|
||||
#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
|
||||
@@ -128,8 +128,8 @@ class LinearFiltration(Strategy):
|
||||
if self.A is None:
|
||||
self._compute_A(
|
||||
model = model,
|
||||
num_classes = num_classes,
|
||||
loader = full_loader,
|
||||
#num_classes = num_classes,
|
||||
loader = forget_loader,
|
||||
device = device
|
||||
)
|
||||
|
||||
@@ -137,7 +137,7 @@ class LinearFiltration(Strategy):
|
||||
Z = self._compute_z(tensor=self.A, forget_index=forget_index)
|
||||
B_Z_rows = []
|
||||
|
||||
for i in range(num_classes):
|
||||
for i in range(self.num_classes):
|
||||
if i == forget_index:
|
||||
B_Z_rows.append(Z)
|
||||
else:
|
||||
@@ -161,8 +161,14 @@ class LinearFiltration(Strategy):
|
||||
|
||||
# overriden function
|
||||
def _split_data(self, dataset):
|
||||
return get_unlearning_loaders(
|
||||
'''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
|
||||
)
|
||||
@@ -2,75 +2,49 @@ import time
|
||||
from pathlib import Path
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.optim as optim
|
||||
from torch.utils.data import DataLoader
|
||||
from unlearning.Strategy import Strategy
|
||||
from architectures.Model import Model, Architecture
|
||||
|
||||
class Retrain(Strategy):
|
||||
"""
|
||||
Implements the Exact Unlearning Baseline by retraining the model architecture
|
||||
completely from scratch using only the retained dataset partition.
|
||||
Implements the Exact Unlearning Baseline by re-instantiating a fresh,
|
||||
pre-trained instance of the specific architecture and training it from scratch
|
||||
on the retain set using the Model's internal train function.
|
||||
"""
|
||||
def __init__(self, target_class_index: int, lr: float = 0.01,
|
||||
weight_decay: float = 0.0005, epochs: int = 5):
|
||||
def __init__(self, target_class_index: int, arch: Architecture, size: int,
|
||||
lr: float = 0.001, epochs: int = 5):
|
||||
super().__init__(target_class_index)
|
||||
self.arch = arch
|
||||
self.size = size
|
||||
self.lr = lr
|
||||
self.weight_decay = weight_decay
|
||||
self.epochs = epochs
|
||||
|
||||
def _reset_weights(self, model: nn.Module):
|
||||
"""
|
||||
Re-initializes all learnable parameters of the model to clear pre-trained memories.
|
||||
"""
|
||||
inner_model = getattr(model, "model", model)
|
||||
for layer in inner_model.modules():
|
||||
if hasattr(layer, 'reset_parameters'):
|
||||
layer.reset_parameters()
|
||||
|
||||
def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
|
||||
# 1. Determine the active execution device from the running sandbox
|
||||
device = next(model.parameters()).device
|
||||
|
||||
print(">> Triggering Exact Unlearning Baseline (Retraining from scratch)...")
|
||||
print(f">> Triggering Exact Unlearning Baseline (Retraining {self.arch.name} from pristine state)...")
|
||||
|
||||
# 1. Clear the pre-trained state completely
|
||||
self._reset_weights(model) # model should be loaded here or weights reset to ImageNet (pretrained default)
|
||||
model.train()
|
||||
# a new model with default params is created
|
||||
fresh_meat = Model.create(self.arch, device, self.size)
|
||||
|
||||
# fresh optimizer for this clean environment
|
||||
optimizer = optim.SGD(
|
||||
model.parameters(),
|
||||
lr=self.lr,
|
||||
momentum=0.9,
|
||||
weight_decay=self.weight_decay
|
||||
# we train it with retain set
|
||||
fresh_meat.train(
|
||||
epochs=self.epochs,
|
||||
loader=retain_loader,
|
||||
rate=self.lr,
|
||||
mode="retrain"
|
||||
)
|
||||
criterion = nn.CrossEntropyLoss()
|
||||
|
||||
# 3. Standard training loop over the Retain set
|
||||
for epoch in range(self.epochs):
|
||||
running_loss = 0.0
|
||||
total_samples = 0
|
||||
|
||||
for data, targets in retain_loader:
|
||||
data, targets = data.to(device), targets.to(device)
|
||||
|
||||
optimizer.zero_grad()
|
||||
outputs = model(data)
|
||||
loss = criterion(outputs, targets)
|
||||
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
running_loss += loss.item() * targets.size(0)
|
||||
total_samples += targets.size(0)
|
||||
|
||||
epoch_loss = running_loss / total_samples
|
||||
print(f" [Retrain] Epoch {epoch+1}/{self.epochs} completed. Loss: {epoch_loss:.4f}")
|
||||
|
||||
print(">> Retraining pipeline finished. Baseline weights established.")
|
||||
|
||||
# 4. Extract the trained nn.Module parameter state dict
|
||||
# In-place copy onto the existing sandbox model structure to seamlessly retain downstream evaluations
|
||||
model.load_state_dict(fresh_meat.model.state_dict())
|
||||
|
||||
print(">> Retraining pipeline finished. Pristine baseline weights successfully established.")
|
||||
return model
|
||||
|
||||
def _split_data(self, dataset):
|
||||
# Dynamically pulls loaders from your Data.py script
|
||||
from sets.Data import get_unlearning_loaders
|
||||
return get_unlearning_loaders(
|
||||
dataset=dataset,
|
||||
|
||||
Reference in New Issue
Block a user