optimised

sets/Data.py

@@ -0,0 +1,239 @@
+from torchvision import datasets, transforms
+from torch.utils.data import Dataset, DataLoader, Subset, ConcatDataset
+import torch
+import numpy as np
+import os
+
+from enum import Enum, auto
+class Set_Name(Enum):
+    CELEBA = auto()
+    CASIAFACES = auto()
+
+# train set transform 
+def train_transform(res):
+    return transforms.Compose([
+        # ResNet expects 224 x 224 res 
+        # Inception expects 299 x 299
+        transforms.Resize((res, res)),
+        transforms.RandomHorizontalFlip(p=0.5),
+        transforms.ColorJitter(
+            brightness=0.2, 
+            contrast=0.2, 
+            saturation=0.1
+        ),
+        transforms.ToTensor(),
+        # normalise to 
+        transforms.Normalize(
+            mean=[0.485, 0.456, 0.406], 
+            std=[0.229, 0.224, 0.225]
+        )
+    ])
+
+# test set transform
+def test_transform(res):
+    return transforms.Compose([
+        transforms.Resize((res, res)),
+        transforms.ToTensor(),
+        transforms.Normalize(
+            mean=[0.485, 0.456, 0.406], 
+            std=[0.229, 0.224, 0.225]
+        )
+    ])
+
+# Load data with 'identity' as target and transform it
+def get_set(set_name:Set_Name):
+    return fetch_celeb_a() if set_name == Set_Name.CELEBA else fetch_casia_faces()
+
+def fetch_celeb_a():
+    return datasets.CelebA(
+        root='./data', 
+        split='all', 
+        target_type='identity', 
+        download=True, 
+        transform=None
+    )
+
+def fetch_casia_faces():
+    # location of the data (path relative to project root)
+    final_path = os.path.abspath("./data/casia-set")
+    
+    if not os.path.exists(final_path):
+        raise FileNotFoundError(
+            f"Unpacked dataset not found at {final_path}. "
+            "Please run Extractor.py first!"
+        )
+    
+    print(f"Loading unpacked CASIA dataset from: {final_path}")
+    
+    return datasets.ImageFolder(
+        root=final_path, 
+        transform=None
+    )
+
+
+
+def get_ids_and_counts(dataset):
+
+    target = get_target(dataset=dataset)
+    return  torch.unique(
+        input = target,
+        return_counts = True
+    )
+
+
+
+# filter selected identities from dataset
+# How many classes, how many images per class
+def select_ids( dataset, sample_size, class_size):
+    ids, counts = get_ids_and_counts(dataset = dataset)
+    eligible_mask = counts >= sample_size
+    eligible_ids = ids[eligible_mask].numpy()
+
+    if len(eligible_ids) < class_size:
+        raise ValueError(
+            f"Only found {len(eligible_ids)} identities with {sample_size}+ images."
+        )
+
+    # Randomly select identities
+    return np.random.choice(eligible_ids, class_size, replace=False)
+
+# optional function to get max amount of samples per class
+def select_top_ids(dataset, class_size):
+    ids, counts = get_ids_and_counts(dataset = dataset)
+
+    # sort by number of images (descending)
+    sorted_indices = torch.argsort(counts, descending = True)
+
+    top_ids = ids[sorted_indices][:class_size].numpy()
+
+    return np.array(top_ids, dtype=int)
+
+
+def get_target(dataset):
+    """
+    Unified target extractor. 
+    Instantly reads raw dataset arrays or safely scales down to unpack wrapped Subsets.
+    """
+    if hasattr(dataset, 'identity'):
+        # celebA
+        targets = dataset.identity
+    elif hasattr(dataset, 'targets'):
+        # others
+        targets = dataset.targets
+    else:
+        # If it's an IdentitySubset or standard Subset, extract mapped targets sequentially
+        # This guarantees we get the 0 -> (n-1) remapped labels
+        targets = [dataset[i][1] for i in range(len(dataset))]
+
+    if not isinstance(targets, torch.Tensor):
+        targets = torch.tensor(targets)
+        
+    return targets
+
+
+# split class images to train and test set.
+def get_indices(dataset, identities, split_at, size = 30):
+
+    if split_at >= size: # debug safety
+        raise ValueError(f"Split point ({split_at}) must be less than total size ({size}).")
+    
+    train_indices = []
+    test_indices = []
+
+    target = get_target(dataset=dataset)
+
+    #training_sample = int(sample_size * training_ratio)
+    np.random.seed(42)
+    for person_id in identities:
+        # Get all indices for this specific person
+        indices = torch.where(target == person_id)[0].numpy()
+
+        # Shuffle the indices for this person
+        np.random.shuffle(indices)
+        
+        # split data to testing and training
+        train_indices.extend(indices[:split_at])
+        test_indices.extend(indices[split_at:size])
+
+    return train_indices, test_indices
+
+
+
+def get_unlearning_loaders(dataset: Dataset, forget_class_idx: int, batch_size: int = 32) -> tuple[DataLoader, DataLoader]:
+    """
+    Splits an IdentitySubset or standard Dataset into forget and retain sets 
+    based on a remapped target class index.
+    """
+    # extract targets
+    targets = get_target(dataset=dataset)
+        
+    # mask indices local to this subset
+    forget_indices = torch.where(targets == forget_class_idx)[0].tolist()
+    retain_indices = torch.where(targets != forget_class_idx)[0].tolist()
+    
+    # PyTorch Subsets
+    forget_subset = Subset(dataset, forget_indices)
+    retain_subset = Subset(dataset, retain_indices)
+    
+    # DataLoaders
+    forget_loader = DataLoader(forget_subset, batch_size=batch_size, shuffle=False)
+    retain_loader = DataLoader(retain_subset, batch_size=batch_size, shuffle=True)
+    
+    print(f"[Data Split] Local Class {forget_class_idx}: {len(forget_subset)} samples | Remaining Classes: {len(retain_subset)} samples.")
+    
+    return retain_loader, forget_loader
+
+
+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:
+        50% of each class goes to the Retain Set, 50% goes to the Forget Set.
+        """
+        
+        # 1. Group dataset indices by their respective ground-truth classes
+        class_to_indices = {c: [] for c in range(num_classes)}
+        
+        print("  [Vertical Split] Tracking class indices across the combined dataset...")
+        for idx in range(len(dataset)):
+            # Extract the label cleanly from the underlying dataset structure
+            _, label = dataset[idx]
+            if label in class_to_indices:
+                class_to_indices[label].append(idx)
+
+        retain_indices = []
+        forget_indices = []
+
+        # 2. Slice each class identity vertically (exactly 50/50)
+        for c, indices in class_to_indices.items():
+            if len(indices) < 2:
+                print(f"  Warning: Class {c} has fewer than 2 samples. Cannot split vertically.")
+                retain_indices.extend(indices)
+                continue
+                
+            # Deterministic shuffle per class to ensure honest distribution before splitting
+            np.random.shuffle(indices)
+            
+            mid = len(indices) // 2
+            forget_indices.extend(indices[:mid])   # First half assigned to unlearning
+            retain_indices.extend(indices[mid:])   # Second half assigned to retention
+
+        print(f"  Vertical split complete: Retain Index Size = {len(retain_indices)} | Forget Index Size = {len(forget_indices)}")
+
+        # 3. Construct lightweight PyTorch Subsets using our sliced index maps
+        retain_subset = Subset(dataset, retain_indices)
+        forget_subset = Subset(dataset, forget_indices)
+
+        # 4. Return pristine, shuffled DataLoaders mirroring your environment's batch specifications
+        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
+
+def _combine_set(loader_one, loader_two):
+    full_train_dataset = ConcatDataset([loader_one.dataset, loader_two.dataset])
+    return DataLoader(
+        full_train_dataset, 
+        batch_size=loader_one.batch_size, 
+        shuffle=True
+    )