reports from optimised linear filtration

sets/Data.py

@@ -186,7 +186,7 @@ def get_unlearning_loaders(dataset: Dataset, forget_class_idx: int, batch_size:
     return retain_loader, forget_loader
 
 
-def vertical_split(dataset, batch_size,num_classes):
+'''def vertical_split(dataset, batch_size,num_classes):
         """
         Executes a class-wise vertical split. 
         Divides the samples of every single identity class exactly in half:
@@ -229,7 +229,61 @@ def vertical_split(dataset, batch_size,num_classes):
         retain_loader = DataLoader(retain_subset, batch_size=batch_size, shuffle=True)
         forget_loader = DataLoader(forget_subset, batch_size=batch_size, shuffle=True)
 
-        return retain_loader, forget_loader
+        return retain_loader, forget_loader'''
+
+def vertical_split(dataset, batch_size, num_classes, ratio = 0.5):
+    """
+    Optimized class-wise vertical split. 
+    Divides the samples of every single identity class exactly in half.
+    """
+    class_to_indices = {c: [] for c in range(num_classes)}
+    
+    # Read cached targets directly instead of decoding image files
+    if hasattr(dataset, 'targets'):
+        labels = dataset.targets
+    elif hasattr(dataset, 'labels'):
+        labels = dataset.labels
+    else:
+        # Fallback only if no cached targets attribute exists
+        labels = [dataset[i][1] for i in range(len(dataset))]
+
+    print("  [Vertical Split] Fast-tracking class indices map...")
+    for idx, label in enumerate(labels):
+        if label in class_to_indices:
+            class_to_indices[label].append(idx)
+
+    retain_indices = []
+    forget_indices = []
+
+    for _, indices in class_to_indices.items():
+        
+        np.random.shuffle(indices) 
+        
+        mid = int(len(indices) * ratio)
+
+        if mid == 0:
+            retain_indices.extend(indices)
+            continue
+
+        forget_indices.extend(indices[:mid])   
+        retain_indices.extend(indices[mid:])   
+
+    print(f"  Vertical split complete: Retain = {len(retain_indices)} | Forget = {len(forget_indices)}")
+
+    retain_subset = Subset(dataset, retain_indices)
+    forget_subset = Subset(dataset, forget_indices)
+
+    # Performance flags keep workers alive between loop runs
+    retain_loader = DataLoader(
+        retain_subset, batch_size=batch_size, shuffle=True, 
+        num_workers=2, pin_memory=True, persistent_workers=True
+    )
+    forget_loader = DataLoader(
+        forget_subset, batch_size=batch_size, shuffle=True, 
+        num_workers=2, pin_memory=True, persistent_workers=True
+    )
+
+    return retain_loader, forget_loader
 
 def _combine_set(loader_one, loader_two):
     full_train_dataset = ConcatDataset([loader_one.dataset, loader_two.dataset])