attck metrics

unlearning/Retrain.py

@@ -1,4 +1,5 @@
-import time
+
+import os
 from pathlib import Path
 import torch
 import torch.nn as nn
@@ -21,27 +22,51 @@ class Retrain(Strategy):
         self.epochs = epochs
 
     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
+
+        # 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">> Triggering Exact Unlearning Baseline (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_meat = Model.create(self.arch, device, self.size)
+        fresh = Model.create(self.arch, device, total_classes)
         
         # we train it with retain set
-        fresh_meat.train(
+        fresh.train(
             epochs=self.epochs, 
             loader=retain_loader, 
             rate=self.lr, 
             mode="retrain"
         )
         
-        # 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())
+        # Extract module parameter state dict and copy in place
+        model.load_state_dict(fresh.model.state_dict())
         
-        print(">> Retraining pipeline finished. Pristine baseline weights successfully established.")
+        print("Retraining pipeline complete")
         return model
 
     def _split_data(self, dataset):
@@ -49,5 +74,5 @@ class Retrain(Strategy):
         return get_unlearning_loaders(
             dataset=dataset,
             forget_class_idx=self.target_class_index,
-            batch_size=32
+            batch_size=16
         )