added reports and params

unlearning/Retrain.py

@@ -0,0 +1,79 @@
+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
+        )