reports from optimised linear filtration

unlearning/Retrain.py

@@ -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,