78 lines
3.0 KiB
Python
78 lines
3.0 KiB
Python
|
|
import os
|
|
from pathlib import Path
|
|
import torch
|
|
import torch.nn as nn
|
|
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 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, 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.epochs = epochs
|
|
|
|
def _run(self, model: nn.Module, forget_loader: DataLoader, retain_loader: DataLoader) -> nn.Module:
|
|
|
|
device = next(model.parameters()).device
|
|
|
|
# we need to check if a retrained copy exists on disk
|
|
checkpoint_path = f"trained_models/class_{self.target_class_index}_retrained.pth"
|
|
if os.path.exists(checkpoint_path):
|
|
print(f"Found existing retrained model checkpoint at '{checkpoint_path}'. Loading parameters directly...")
|
|
|
|
# Load the state dict using safe configuration flags
|
|
state_dict = torch.load(checkpoint_path, map_location=device, weights_only=True)
|
|
|
|
# Safely apply the parameter weights to the model in-place
|
|
model.load_state_dict(state_dict)
|
|
print("Retrained parameter loading complete (Retraining bypassed).")
|
|
return model
|
|
|
|
# Cache Miss: Execute the standard retraining pipeline
|
|
print(f"No naive model found for class {self.target_class_index} retraining a new one")
|
|
|
|
|
|
print(f"Retraining {self.arch.name} from pristine state)...")
|
|
inner_model = getattr(model, "model", model)
|
|
if hasattr(inner_model, "fc"):
|
|
total_classes = inner_model.fc.out_features
|
|
elif hasattr(inner_model, "classifier"):
|
|
# Fallback for alternative architecture layout types
|
|
total_classes = inner_model.classifier[-1].out_features
|
|
else:
|
|
total_classes = self.size
|
|
|
|
# a new model with default params is created
|
|
fresh = Model.create(self.arch, device, total_classes)
|
|
|
|
# we train it with retain set
|
|
fresh.train(
|
|
epochs=self.epochs,
|
|
loader=retain_loader,
|
|
rate=self.lr,
|
|
mode="retrain"
|
|
)
|
|
|
|
# Extract module parameter state dict and copy in place
|
|
model.load_state_dict(fresh.model.state_dict())
|
|
|
|
print("Retraining pipeline complete")
|
|
return model
|
|
|
|
def _split_data(self, dataset):
|
|
from sets.Data import get_unlearning_loaders
|
|
return get_unlearning_loaders(
|
|
dataset=dataset,
|
|
forget_class_idx=self.target_class_index,
|
|
batch_size=16
|
|
) |