79 lines
2.9 KiB
Python
79 lines
2.9 KiB
Python
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
|
|
|
|
class Retrain(Strategy):
|
|
"""
|
|
Implements the Exact Unlearning Baseline by retraining the model architecture
|
|
completely from scratch using only the retained dataset partition.
|
|
"""
|
|
def __init__(self, target_class_index: int, lr: float = 0.01,
|
|
weight_decay: float = 0.0005, epochs: int = 5):
|
|
super().__init__(target_class_index)
|
|
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:
|
|
device = next(model.parameters()).device
|
|
|
|
print(">> Triggering Exact Unlearning Baseline (Retraining from scratch)...")
|
|
|
|
# 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()
|
|
|
|
# fresh optimizer for this clean environment
|
|
optimizer = optim.SGD(
|
|
model.parameters(),
|
|
lr=self.lr,
|
|
momentum=0.9,
|
|
weight_decay=self.weight_decay
|
|
)
|
|
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.")
|
|
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,
|
|
forget_class_idx=self.target_class_index,
|
|
batch_size=32
|
|
) |